kyegomez / open-jev

# Open Jev
An open-source, from-first-principles reconstruction of the ideas behind
TypeSafe AI's Jev,
written in PyTorch.
> [!IMPORTANT]
> This is an unofficial research implementation with random weights. It is not
> the production Jev model, does not reproduce TypeSafe AI's training data or
> weights, and makes no claim of matching their published results.
## The Idea
Jev is presented as a system one model: unstructured program state goes in,
typed probabilistic decisions come out.
Instead of generating text one token at a time, this implementation encodes the
state once and answers every question through small typed readout heads.
 [code] 
Questions are isolated from one another and folded into the batch dimension.
They may attend to the shared state, but not to other questions. The result is a
single parallel forward pass with outputs constrained by their declared types.
## Architecture
- A bidirectional transformer encodes nested state without a causal mask
- Structural path embeddings preserve where values occur in dictionaries and arrays
- Learned query slots cross-attend into one shared, cacheable state representation
- `Noul` returns the probability that a statement is true
- `Choice` softmaxes only over the options supplied at runtime
- `Score` uses ordinal thresholds and returns a distribution plus its expectation
- An evidential head models epistemic confidence separately from class probability
- `RLCDLoss` combines soft-target NLL, Brier score, consistency, confidence, and ECE terms
The implementation is intentionally small and legible. The included tokenizer
is a deterministic hash tokenizer, suitable for running the architecture but
not for training a useful model.
## Install
 [code] 
Python 3.10 or newer is recommended.
## Usage
 [code] 
The model is untrained, so the values from this example are random. The useful
guarantee is structural: a `Choice` answer can only be one of the options that
were declared.
To run the complete demo, including one calibration-oriented training step:
 [code] 
## Training
`RLCDLoss` expects probability distributions rather than one-hot labels. This
allows disagreement and ambiguity to remain visible instead of forcing every
example toward certainty.
 [code] 
For consistency training, pass semantically equivalent states through
`augmented_states`—for example, paraphrases or shuffled dictionary keys.
## Notes
This repository explores an architecture inferred from public material. Details
such as the encoder design, query slots, ordinal head, confidence objective, and
training recipe are hypotheses, not disclosed details of TypeSafe AI's system.
See `docs/ARCHITECTURE.md` for the full reasoning and
tradeoffs behind the reconstruction.
In particular, typed output prevents schema violations; it does not guarantee
that a prediction is correct. Calibration must also be measured again under
distribution shift before deploying a model in an unattended workflow.
## Todo
- [ ] Replace the hash tokenizer with a