openjev-sdk 0.2.0__tar.gz

openjev_sdk-0.2.0/.gitignore

@@ -0,0 +1,46 @@
+.venv/
+models/
+__pycache__/
+*.pyc
+.DS_Store
+runs/
+*.npz
+# ViZDoom writes this config into the cwd on first launch
+_vizdoom.ini
+
+# mkdocs build output
+site/
+adapters/
+data/chess/
+data/wikiracing/
+
+# Local planning notes
+/ideas/
+
+# Built Python distributions
+/dist/
+
+# Claude Code local settings and session files
+.claude/
+CLAUDE.local.md
+
+# Temp, editor and tool files
+*.tmp
+*.temp
+*.swp
+*.swo
+*~
+*.log
+.idea/
+.vscode/
+.pytest_cache/
+.mypy_cache/
+.ruff_cache/
+.ipynb_checkpoints/
+.coverage
+htmlcov/
+*.egg-info/
+build/
+.env
+.env.*
+graphify-out/

openjev_sdk-0.2.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Dasein Labs
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

openjev_sdk-0.2.0/PKG-INFO

@@ -0,0 +1,367 @@
+Metadata-Version: 2.5
+Name: openjev-sdk
+Version: 0.2.0
+Summary: Train small task-oriented decision models on your own traces: pick, calibrate, defer, and improve with RL
+Project-URL: Homepage, https://github.com/daseinlabs/open-jev
+Project-URL: Documentation, https://github.com/daseinlabs/open-jev/blob/main/docs/tod-sdk.md
+Project-URL: Cookbooks, https://github.com/daseinlabs/open-jev/tree/main/cookbooks
+Author: Dasein Labs
+License-Expression: MIT
+License-File: LICENSE
+Keywords: agents,cuda,decision model,fine-tuning,gemma,mlx,pytorch,reinforcement learning,tool calling
+Classifier: Development Status :: 3 - Alpha
+Classifier: Environment :: GPU :: NVIDIA CUDA
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.12
+Requires-Dist: fastapi>=0.141.1
+Requires-Dist: huggingface-hub>=1.31.0
+Requires-Dist: mlx-lm>=0.24; sys_platform == 'darwin' and platform_machine == 'arm64'
+Requires-Dist: mlx>=0.24; sys_platform == 'darwin' and platform_machine == 'arm64'
+Requires-Dist: numpy>=2.5.3
+Requires-Dist: safetensors>=0.8.0
+Requires-Dist: torch>=2.6; sys_platform != 'darwin' or platform_machine != 'arm64'
+Requires-Dist: transformers>=5.17.0
+Requires-Dist: uvicorn[standard]>=0.53.0
+Provides-Extra: doom
+Requires-Dist: vizdoom>=1.2.4; extra == 'doom'
+Provides-Extra: finetune
+Requires-Dist: chess>=1.10; extra == 'finetune'
+Requires-Dist: pyarrow>=15; extra == 'finetune'
+Provides-Extra: torch
+Requires-Dist: accelerate>=1.0; extra == 'torch'
+Requires-Dist: torch>=2.6; extra == 'torch'
+Description-Content-Type: text/markdown
+
+# openjev
+
+One-pass option scoring with a local Gemma 3 4B: MLX on Apple silicon, or PyTorch on Windows and Linux (CPU or GPU).
+Design notes: [docs/design/one-pass-option-scoring.md](docs/design/one-pass-option-scoring.md);
+per-task training: [docs/design/per-task-finetuning-with-gemma.md](docs/design/per-task-finetuning-with-gemma.md).
+
+Given a context and a list of pre-written options, the model prefills the
+context once, expands that KV cache across the option batch, and scores every
+option in a single padded forward pass. No decoding. The score is the
+log-probability of the option tokens given the context; a softmax over the
+option scores gives a probability per option, like `jevlike-predict`.
+
+## Train a decision model on your own data
+
+`openjev.tod` trains small task-oriented decision models: a frozen Gemma backbone plus a
+~2M-parameter head that picks one option from a list your code supplies (which tool,
+which order, which team), reports a calibrated confidence, and defers to a bigger model
+when unsure. Runs on NVIDIA/AMD GPUs (PyTorch) and Apple silicon (MLX). Install with
+`pip install openjev-sdk`.
+
+```python
+from openjev.tod import Candidate, DecisionSpec, Episode, Kind, build_dataset
+
+TEAMS = {'billing': 'Billing and refunds', 'shipping': 'Delivery and tracking'}
+spec = DecisionSpec('routing', [Kind('team', 'Which team should handle this?',
+    lambda state: [Candidate(k, v, payload=k) for k, v in TEAMS.items()], version='1')])
+episodes = [Episode(t['id'], (spec.decision(t['id'], 'team', {'message': t['text']}, choice=t['team']),))
+            for t in tickets]
+build_dataset(spec, episodes).write('runs/routing')
+```
+
+```bash
+openjev tod train runs/routing --model models/gemma-3-4b-it --out runs/routing-model
+```
+
+Chat traces, RL in a simulator, and deriving the spec from your tool schemas are covered in
+the [SDK guide](https://github.com/daseinlabs/open-jev/blob/main/docs/tod-sdk.md) and the three [cookbooks](https://github.com/daseinlabs/open-jev/tree/main/cookbooks).
+
+## Demo: Doom in the terminal
+
+`demo/doom/` runs ViZDoom headless, describes each frame in a line of text, and
+lets the server rank the action menu with one `/score` call (or one System One
+`choice` question). Start `make serve`, then `make doom`. Details and keys in
+[`demo/doom/README.md`](demo/doom/README.md).
+
+![openjev playing Doom in the terminal: the model ranks the action menu each step](https://raw.githubusercontent.com/daseinlabs/open-jev/main/docs/media/doom-recording.gif)
+
+Full-resolution recording: [docs/media/doom-recording.mov](docs/media/doom-recording.mov).
+
+## Setup
+
+For use inside a Python application or notebook, install this checkout with
+`python -m pip install .` (Python 3.12+). See the [Python package guide](docs/python-package.md)
+for decision inference, feature extraction, training, and wheel installation.
+
+### Windows and Linux
+
+Install Python 3.12+ and [uv](https://docs.astral.sh/uv/getting-started/installation/), then run these commands from the repository in PowerShell or a Linux shell:
+
+```sh
+uv sync
+uv run hf auth login
+uv run hf download google/gemma-3-4b-it --local-dir models/gemma-3-4b-it
+uv run openjev serve --backend torch --device auto --port 8000
+```
+
+Accept the Gemma license on Hugging Face before downloading. You can also pass
+`--model google/gemma-3-4b-it` to load directly from the Hugging Face cache.
+No Make, Xcode, shell activation, or `.venv/bin` paths are needed here.
+
+`--backend auto` (the default) selects MLX on Apple silicon and PyTorch elsewhere.
+For PyTorch, `--device auto` selects CUDA when available, then MPS, then CPU.
+Use `--device cuda` or `--device cuda:1` to require a specific GPU (fails clearly
+if CUDA is unavailable), or `--device cpu` to force CPU inference.
+
+For NVIDIA acceleration, install the GPU driver and a matching PyTorch build using
+the [official PyTorch installer](https://pytorch.org/get-started/locally/).
+Run its pip command in this repository's virtual environment (use `uv pip install`
+in place of `pip3 install`). After a custom PyTorch install, use `uv run --no-sync`
+so uv does not replace your selected build. On supported Linux AMD systems,
+PyTorch ROCm builds also use the `cuda` device name.
+
+```sh
+uv run --no-sync python -c "import torch; print(torch.__version__, torch.cuda.is_available())"
+uv run --no-sync openjev serve --backend torch --device cuda --batch-size 2
+uv run --no-sync openjev check --backend torch --device cuda
+uv run --no-sync openjev score --backend torch --device cuda --context "The capital of France is" --option " Paris" --option " Berlin"
+```
+
+The base 4B weights need roughly 8 GB just for GPU model weights at 16-bit precision,
+plus memory for activations and the option KV caches. Reduce `--batch-size` and
+context length if you run out of GPU memory. CPU runs use float32 and need more RAM.
+Use original Hugging Face weights with PyTorch; MLX quantized weights and MLX LoRA
+adapters are not supported by this backend.
+
+Scoring, evaluation, benchmarks, correctness checks, both HTTP scoring endpoints,
+frozen-feature extraction, head training and evaluation, `openjev decision`, and
+`openjev tod` (including RL) support PyTorch. Contextual features (`--contextual`) and the
+chess LoRA training workflow still require MLX on Apple silicon. The Doom terminal UI
+also has its own platform dependencies; cross-platform support here covers the
+scoring CLI and HTTP server.
+
+### Apple silicon / existing Make workflow
+
+```sh
+make setup       # uv sync (arm64 Python 3.12 venv) + download google/gemma-3-4b-it into models/ (gated; needs HF login)
+make serve       # start the HTTP server on :8000
+make health      # curl /health
+make request     # example curl against /score
+make systemone   # TypeSafe-style request against /v1/systemone
+make check       # verify cached batched scoring against naive re-encoding
+make bench       # latency benchmark
+make eval DATA=data/synthetic/test.jsonl
+```
+
+`make setup` also installs the optional `torch` extra, used only for the jevlike comparison and HF cross-checks.
+`make` on macOS needs the Xcode licence accepted (`sudo xcodebuild -license accept`) or Homebrew's `gmake`.
+
+## Usage
+
+```sh
+# Rank options for one context (prints probability, score, raw sum, token count)
+.venv/bin/openjev score --context "The capital of France is" \
+    --option " Paris" --option " Berlin" --option " Lyon"
+
+# Chat template (context as user turn, options scored as the reply) + PMI normalisation
+.venv/bin/openjev score --chat --norm pmi --context "..." --option "..." --option "..."
+
+# Predefined options: one per line in a text file, reused for every context
+.venv/bin/openjev score --options-file options.txt --context "..."
+.venv/bin/openjev eval contexts.jsonl --fixed-options options.txt   # rows need only {"context": ...}; add "label" for accuracy
+
+# Top-1 / top-3 accuracy on jevlike-style JSONL: {"context": ..., "options": [...], "label": 0}
+.venv/bin/openjev eval data.jsonl --norm mean
+
+# Verify the cached batched path against naive re-encoding, and benchmark it
+.venv/bin/openjev check
+.venv/bin/openjev bench --context-tokens 200 --options 8 --option-tokens 30
+```
+
+## Server
+
+Loads the model once. The measured Apple silicon MLX path answers scoring requests in about
+90 ms each; PyTorch latency depends on your CPU/GPU. Use `uv run openjev serve` on any
+supported platform, or the Make shortcut below on macOS.
+
+```sh
+make serve                                     # = .venv/bin/openjev serve --port 8000
+curl -s localhost:8000/score -H 'content-type: application/json' -d '{
+  "context": "Customer: my order arrived broken. Agent:",
+  "options": [" I am sorry, I will send a replacement.", " Please read our returns policy."],
+  "norm": "mean", "chat": false, "sep": ""
+}'
+```
+
+The response has `best`, `best_index`, per-option `probability` / `logprob_sum` / `n_tokens`, and `timing`.
+
+## TypeSafe System One contract
+
+`POST /v1/systemone` implements the request/response shape documented at
+[docs.typesafe.ai](https://docs.typesafe.ai): a `state` (string, object or array) plus a map of typed
+`questions`, answered against that state.
+
+```sh
+make systemone     # sends examples/systemone-quickstart.json; or:
+curl -s localhost:8000/v1/systemone -H 'content-type: application/json' -d '{
+  "state": "I ordered size 10 shoes but received size 8. Please send the right size.",
+  "model": "jev-latest",
+  "questions": {
+    "department":  {"type": "choice", "instructions": "Which department handles this?",
+                    "criteria": {"returns": "Returns and exchanges", "shipping": "Delivery issues", "billing": "Charges and refunds"}},
+    "severity":    {"type": "score",  "instructions": "How severe is the problem?",
+                    "criteria": ["minor", "moderate: wrong item", "major: safety or financial loss"]},
+    "wants_refund":{"type": "noul",   "instructions": "Is the customer asking for a refund to their card?"}
+  }
+}'
+```
+
+| question `type` | request `criteria` | answer fields |
+|---|---|---|
+| `choice` | map of option name to description (string, object, array or null); up to 255 options | `choice`, `probabilities` (sum to 1), `confidence` |
+| `score` | ordered array of level descriptions | `score` (probability-weighted mean of level index), `probabilities` keyed `"0".."n-1"`, `legend`, `confidence` |
+| `noul` | optional `{"true": ..., "false": ...}` | `noul` = probability of yes |
+
+Response: `{"model", "answers": {id: answer}, "usage": {"input_tokens", "output_tokens"}}`.
+Set `OPENJEV_API_KEY` before `make serve` to require `Authorization: Bearer <key>` on this route.
+
+How it maps onto the scorer: each question is rendered to a plain-text prompt (state, instructions,
+options or levels, then `Answer:` and a newline) and the option names, level numbers, or `yes`/`no`
+are scored as continuations in one prefix-shared batched pass per question. `confidence` is 1 minus
+the normalised entropy of the distribution; TypeSafe does not publish its formula, so treat it as an
+approximation. `usage.output_tokens` counts the candidate-label tokens that were scored.
+
+Comparison on the docs' quick-start request (`examples/systemone-quickstart.json`), Gemma 3 4B
+zero-shot vs the numbers TypeSafe publishes for Jev:
+
+| answer | Jev (docs) | openjev / Gemma 3 4B |
+|---|---|---|
+| `department.choice` | technical, p=0.84, confidence 0.60 | technical, p=1.00, confidence 1.00 |
+| `frustration.score` | 1.04 (annoyed but polite) | 2.00 (furious) |
+| `is_urgent.noul` | 0.999 | 0.005 |
+
+Same routing decision, but Gemma is over-confident and disagrees on the two judgement calls. Zero-shot
+probabilities are softmaxed next-token likelihoods, not calibrated judgements. Closing the gap means
+labelled data and a trained head (below).
+
+## Native candidate-decision model
+
+`openjev decision` trains a shared candidate-scoring head on frozen features and
+reloads it for raw state/question/candidate requests. It includes versioned task
+validation, feature and artifact fingerprints, held-out evaluation against
+continuation scoring, and candidate-order/repeated-request checks.
+
+See [the native decision-head guide](docs/decision-head.md) for data formats and
+commands. It runs on MLX or PyTorch (`OPENJEV_BACKEND=mlx|torch`, `OPENJEV_DEVICE=cuda|mps|cpu`).
+
+## Training a head (per-task, on frozen Gemma features)
+
+```sh
+make features    # one prefix-shared pass per example, cached to runs/feats/{train,validation,test}.npz
+make train       # jevlike's cross-attention head (MLX or PyTorch), listwise cross-entropy -> runs/head.safetensors
+make eval-head   # top-1/top-3, ECE, and the shuffled-context control on the test split
+```
+
+- `openjev features DATA --out F.npz` stops Gemma before the LM head, keeps every context token's
+  final hidden state and the masked-mean of each option's tokens (float16). Options are encoded on
+  their own, as in jevlike, so the head has to do the matching. `--contextual` encodes them as
+  continuations of the context instead: stronger features, but the match leaks into the option
+  vectors and the shuffled-context control below stops meaning anything.
+- `openjev train` = AdamW 5e-4 (2e-3 diverges on Gemma features, whose norms are ~115), weight decay 1e-4, grad clip 1.0, 8 epochs, batch 64, best validation
+  epoch kept. The checkpoint is `head.safetensors` plus `head.json` (rank, hidden size, training config).
+- `openjev eval-head` reports what `jevlike-eval` reports: top-1, top-3, 10-bin expected calibration
+  error, and the same metrics with every example paired with another example's context. If the
+  shuffled number does not collapse, the head is reading option priors rather than the state.
+
+Result on the synthetic split (2000 train / 400 validation / 400 test, 2026-09-17): features at
+77 ms per example, head training 8 epochs in about 15 s.
+
+| | top-1 | top-3 | ECE |
+|---|---|---|---|
+| head on frozen Gemma features, test | 0.970 | 1.000 | 0.027 |
+| same head, shuffled-context control | 0.258 | 0.698 | 0.724 |
+| jevlike tiny byte encoder, test | 0.998 | 1.000 | 0.008 |
+| Gemma zero-shot, `--norm sum`, test | 1.000 | 1.000 | n/a |
+
+The control sits at chance (about 4.5 options per row), so the head really matches options to the
+state, and its ECE of 0.03 is what a calibrated `confidence` looks like.
+
+Python:
+
+```python
+from openjev import OptionScorer
+s = OptionScorer("models/gemma-3-4b-it", batch_size=8)
+for r in s.score("The capital of France is", [" Paris", " Berlin"], norm="mean"):
+    print(r.option, r.probability, r.logprob_sum, r.n_tokens)
+print(s.last_timing)
+```
+
+### Normalisation (`--norm`)
+
+| norm | score | use when |
+|---|---|---|
+| `mean` (default) | sum of option-token log-probs divided by token count | options differ in length |
+| `sum` | total log-prob | options are the same length or you want raw likelihood |
+| `pmi` | sum minus the option's unconditional log-prob (BOS-only context) | options differ in base-rate plausibility; costs one extra batched pass |
+
+## Measured on M5 Pro, 64 GB (bf16, mlx-lm)
+
+Workload: 202-token context, 8 options, 242 option tokens total.
+
+| path | median latency |
+|---|---|
+| context cached once, options batched | 0.17 s |
+| context re-encoded per option, no cache | 0.68 s |
+
+Model load is about 1 s from a warm disk. First forward pass adds under a second of warm-up.
+
+## Validating against jevlike
+
+`jevlike` is installed into the same venv (`uv pip install -e ../../vinnylarouge/jevlike`), so both
+CLIs read the same JSONL. Generate its synthetic menu set, then score it both ways:
+
+```sh
+.venv/bin/jevlike-data synthetic --output data/synthetic
+.venv/bin/openjev eval data/synthetic/test.jsonl --norm sum --sep $'\nChoice: '
+.venv/bin/jevlike-train data/synthetic/train.jsonl --validation data/synthetic/validation.jsonl \
+    --output runs/synthetic-tiny.pt --device mps
+.venv/bin/jevlike-eval runs/synthetic-tiny.pt data/synthetic/test.jsonl --device mps
+```
+
+Results on the 400-row synthetic test set (2026-09-16):
+
+| scorer | training | top-1 | top-3 | median latency / example |
+|---|---|---|---|---|
+| openjev, Gemma 3 4B zero-shot, `--norm sum` | none | 1.000 | 1.000 | 0.086 s |
+| openjev, `--norm mean` | none | 0.988 | 1.000 | 0.086 s |
+| openjev, `--norm pmi` | none | 0.988 | 1.000 | 0.153 s |
+| jevlike tiny byte encoder + head | 2000 rows, 8 epochs | 0.998 | 1.000 | well under 10 ms |
+
+The synthetic task is easy for both. The real validation is your own labelled rows: run
+`openjev eval` on them zero-shot and compare against a `jevlike-train`/`jevlike-eval` run on the
+same split. If Gemma zero-shot is close to the trained head, Route B is enough; if not, train a head
+(Route A) with `make features && make train && make eval-head`.
+
+## Layout
+
+- `openjev/torch_backend.py`: PyTorch CPU/GPU inference and shared-prefix KV cache scoring.
+- `openjev/scorer.py`: `OptionScorer` (prefill, cache expansion, batched scoring, naive reference).
+- `openjev/systemone.py`: System One request/answer models, prompt renderers, zero-shot answers.
+- `openjev/server.py`: FastAPI app, `/health`, `/score`, `/v1/systemone`.
+- `openjev/features.py`: frozen-Gemma feature extraction and the `.npz` feature cache.
+- `openjev/compute.py`: backend selection (MLX or PyTorch) for features, heads and training.
+- `openjev/head.py`, `openjev/head_torch.py`: jevlike's cross-attention head in `mlx.nn` and `torch.nn`; the saved weights load in either.
+- `openjev/train.py`: head training loop and evaluation (top-k, ECE, shuffled-context control).
+- `openjev/cli.py`: `openjev score | eval | bench | check | serve | features | train | eval-head`.
+- `demo/doom/`: Doom in the terminal, the server picks every action (`make doom`).
+- `models/`: downloaded weights (git-ignored).
+
+## Development checks
+
+```sh
+uv sync --extra torch --group test
+uv run python -m unittest discover -s tests -v
+```
+
+The offline tests use a tiny randomly initialized Gemma model to compare cached,
+padded scoring with full re-encoding, including sliding-window attention, multiple
+batch sizes, normalization, and repeated requests. CI runs them on Windows, Linux,
+and macOS. Actual GPU availability and throughput must be checked on your hardware.