tevro 0.1.0__tar.gz

tevro-0.1.0/.gitignore

@@ -0,0 +1,17 @@
+# Python-generated files
+__pycache__/
+*.py[oc]
+build/
+dist/
+*.egg-info
+
+# Virtual environments
+.venv
+
+# Training runs (LoRA + head checkpoints, metrics)
+runs/
+
+# Local environments and data
+.env
+data/
+.DS_Store

tevro-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Kartik Gadad
+
+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.

tevro-0.1.0/PKG-INFO

@@ -0,0 +1,199 @@
+Metadata-Version: 2.5
+Name: tevro
+Version: 0.1.0
+Summary: A small decision model: picks one option and gives a probability for every option, in one forward pass
+Project-URL: Homepage, https://github.com/kartikg-gk/tevro
+Project-URL: Model, https://huggingface.co/kartikg1/tevro-decider
+Author: Kartik Gadad
+License-Expression: MIT
+License-File: LICENSE
+Keywords: calibration,classification,decision model,lora,qwen3
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.10
+Requires-Dist: huggingface-hub>=0.25.0
+Requires-Dist: peft>=0.21.0
+Requires-Dist: safetensors>=0.6.0
+Requires-Dist: torch>=2.14.0
+Requires-Dist: transformers>=5.17.0
+Description-Content-Type: text/markdown
+
+# Tevro
+
+Tevro is a small decision model, and this repository holds its training code and a pip package for running it. It reads a state, a question and a fixed set of options, and returns one probability per option in a single forward pass. It handles three question types: multiple choice, yes/no, and ordinal scores.
+
+![Tevro architecture](https://raw.githubusercontent.com/kartikg-gk/tevro/main/arch.png)
+
+## Use the model
+
+The trained model is on Hugging Face: **[kartikg1/tevro-decider](https://huggingface.co/kartikg1/tevro-decider)**. It needs Python 3.12 and about 3 GB of free memory; a GPU is optional.
+
+```bash
+pip install tevro
+tevro predict request.json --device cpu
+```
+
+`tevro predict` reads one request from a JSON file (examples below) and prints the answer. The weights and the base model are downloaded on the first run. Use `--device cuda` on an NVIDIA GPU. From Python:
+
+```python
+import tevro
+
+model = tevro.load()                      # downloads kartikg1/tevro-decider on first use
+model.answer("Parcel A is on shelf B2. Parcel B is on shelf C3.",
+             {"type": "choice", "instructions": "Which shelf holds Parcel A?",
+              "criteria": {"a": "B2", "b": "C3"}})
+```
+
+The pip package only runs the model. Training code, data tools and evaluation scripts live in this repository.
+
+**Multiple choice**
+
+```json
+{"state": "Parcel A is on shelf B2. Parcel B is on shelf C3.",
+ "question": {"type": "choice", "instructions": "Which shelf holds Parcel A?",
+              "criteria": {"a": "B2", "b": "C3"}}}
+```
+```json
+{"confidence": 0.99996, "abstain": false, "choice": "a",
+ "probabilities": {"a": 0.99996, "b": 0.00004}}
+```
+
+**Yes/no**
+
+```json
+{"state": "Refunds need a receipt and a purchase within 30 days. The customer bought it 12 days ago and has no receipt.",
+ "question": {"type": "noul", "instructions": "Should the refund be issued?"}}
+```
+```json
+{"confidence": 0.933, "abstain": false, "noul": 0.067}
+```
+
+`noul` is the probability of yes.
+
+**Score**
+
+```json
+{"state": "The app crashes for every user when they open the payments page. No data is lost.",
+ "question": {"type": "score", "instructions": "How severe is this bug?",
+              "criteria": ["Cosmetic only", "Minor, with a workaround", "A core feature is blocked", "Data loss or harm"]}}
+```
+```json
+{"confidence": 0.431, "abstain": false, "score": 1.65,
+ "probabilities": [0.019, 0.431, 0.428, 0.122]}
+```
+
+`score` is the expected level (0 to 3 here). A low `confidence` like this one means the model is unsure; send those cases to a person.
+
+## Results
+
+| Evaluation | Accuracy |
+| --- | ---: |
+| Separate held-out test split (5,000) | 80.2% |
+| — multiple choice / yes-no / score | 80.8% / 88.0% / 65.8% |
+| Training dataset's test split (5,000) | 80.9% |
+| JevBench public (231) | 57.6% |
+| — easy / standard / hard | 97.9% / 66.7% / 34.2% |
+
+On the held-out split, yes/no AUC is 0.953 and expected calibration error is 0.037. At confidence 0.95 or higher the model answers 48.9% of questions and is right on 98.0% of them.
+
+## Set up for training
+
+```bash
+git clone https://github.com/kartikg-gk/tevro
+cd tevro
+uv sync
+uv run pytest -q tests
+```
+
+Python 3.12, managed with [uv](https://docs.astral.sh/uv/). On Windows, `uv sync` installs the CUDA 13.0 build of PyTorch; on Linux it uses the default PyPI build. The tests use a tiny randomly initialised model and run on CPU.
+
+## Train
+
+```bash
+uv run python train.py configs/smoke.yaml
+uv run python train.py configs/check_20k.yaml
+uv run python train.py configs/release_12k.yaml
+```
+
+| Config | What it does | Hardware |
+| --- | --- | --- |
+| `smoke.yaml` | Overfits 10 questions; train accuracy should reach 1.0 | any GPU |
+| `check_20k.yaml` | 20,000 questions, 1,000 test, bf16 weights, micro-batch 2 × 32 | 6 GB laptop GPU, ~45 min |
+| `release_12k.yaml` | The released model: 12,000 updates of 64 questions, fp32 | A100 80 GB, ~6.5 h |
+| `finetune_rules.yaml`, `finetune_v2.yaml` | Short fine-tunes of the released model on a 20K rule-reasoning mix (`--resume` the release checkpoint); ~10 min | A100 80 GB |
+
+Each run writes `runs/<name>-<timestamp>/` with the config, `metrics.jsonl`, `summary.json` and checkpoints (`best`, `latest`, `final`; LoRA adapter and head only). Every training log line shows `eta_min`, the projected minutes left. `--resume <run or checkpoint>` reloads the weights and starts a fresh optimizer.
+
+### On Modal
+
+```bash
+modal run --detach modal_app.py --config configs/release_12k.yaml
+modal run modal_app.py::evaluate --checkpoint runs/<run>/checkpoints/final --out eval.txt
+```
+
+`modal_app.py` uploads this folder, keeps `runs/` and the Hugging Face cache on Modal Volumes, runs on an A100 80 GB (set `TEVRO_GPU` to change it), and caps a training job at 9 hours.
+
+## Evaluate
+
+```bash
+uv run python inference.py runs/<run>/checkpoints/final --max_questions 5000
+uv run python jevbench_eval.py runs/<run>/checkpoints/final <jevbench>/datasets/public
+```
+
+The first scores a checkpoint on the training dataset's test split: accuracy per question type, yes/no AUC, ECE, Brier score, reliability and resolution, and accuracy at confidence 0.9 and 0.95. The second scores the 231 public [JevBench](https://github.com/fstandhartinger/jevbench) items by tier.
+
+## Data
+
+Each row is one state with one or more questions:
+
+```json
+{
+  "state": "Parcel A is on shelf B2. Parcel B is on shelf C3.",
+  "questions_json": "{\"q1\": {\"type\": \"choice\", \"instructions\": \"Which shelf holds Parcel A?\", \"criteria\": {\"a\": \"B2\", \"b\": \"C3\"}, \"label\": \"a\"}}"
+}
+```
+
+- `choice`: `criteria` maps option keys to descriptions; `label` is a key.
+- `noul` (yes/no): `label` is `true` or `false`.
+- `score`: `criteria` is an ordered list of levels; `label` is the level index.
+
+By default the code loads [kartikg1/tevro-decision-v1](https://huggingface.co/datasets/kartikg1/tevro-decision-v1) (250,000 training and 5,000 test questions; private). Point `data.question_files` at your own parquet files (`train` and `test`) to train on your data. `audit.py` reports label-position, length and negation shortcuts before you spend GPU time.
+
+## Model
+
+| Part | Setting |
+| --- | --- |
+| Base | Qwen3-0.6B, first 20 of 28 layers |
+| Adapter | LoRA rank 8, alpha 16, on q/k/v/o of layers 6–19 |
+| Input | task hint + question + state, then each option as its own span starting at the same position, then `The answer is:` |
+| Head | question-type embedding on every token, 2 bidirectional transformer layers (1024-d, 16 heads, dropout 0.1), scalar scorer at each option's last token |
+| Loss | cross-entropy over the options |
+| Optimizer | AdamW, LoRA lr 1e-4, head lr 1e-3, weight decay 0.01, cosine schedule, 5% warmup, gradient clip 1.0 |
+
+Options cannot attend to each other, so reordering them only reorders the probabilities. The test suite checks this.
+
+## Repository layout
+
+```
+tevro/              the pip package: copies of network.py and tokenization.py, plus predict.py and cli.py
+tokenization.py     input layout: positions, attention mask, read indices
+network.py          base model, LoRA, decision head
+dataset.py          loading, question expansion, batching
+train.py            training loop
+inference.py        evaluation metrics and single answers
+calibrate.py        temperature fitting, abstention thresholds, calibration metrics
+jevbench_eval.py    JevBench public evaluation
+audit.py            data shortcut checks
+synthetic.py        procedural decision questions
+data_variants.py    paired question variations
+reasoning_data.py   procedural rule-reasoning questions (policies, payouts, counting)
+build_mix.py        builds a 20K fine-tuning mix from generated, public and replayed questions
+llm_items.py        writes hard items with an LLM through OpenRouter and keeps those a blind second pass agrees with (needs `pydantic` and `OPENROUTER_API_KEY`)
+modal_app.py        Modal training and evaluation
+configs/            run settings
+tests/              layout, invariance, loss and data tests
+```
+
+## License
+
+MIT. The base model, Qwen3-0.6B, is released under its own license.

tevro-0.1.0/README.md

@@ -0,0 +1,179 @@
+# Tevro
+
+Tevro is a small decision model, and this repository holds its training code and a pip package for running it. It reads a state, a question and a fixed set of options, and returns one probability per option in a single forward pass. It handles three question types: multiple choice, yes/no, and ordinal scores.
+
+![Tevro architecture](https://raw.githubusercontent.com/kartikg-gk/tevro/main/arch.png)
+
+## Use the model
+
+The trained model is on Hugging Face: **[kartikg1/tevro-decider](https://huggingface.co/kartikg1/tevro-decider)**. It needs Python 3.12 and about 3 GB of free memory; a GPU is optional.
+
+```bash
+pip install tevro
+tevro predict request.json --device cpu
+```
+
+`tevro predict` reads one request from a JSON file (examples below) and prints the answer. The weights and the base model are downloaded on the first run. Use `--device cuda` on an NVIDIA GPU. From Python:
+
+```python
+import tevro
+
+model = tevro.load()                      # downloads kartikg1/tevro-decider on first use
+model.answer("Parcel A is on shelf B2. Parcel B is on shelf C3.",
+             {"type": "choice", "instructions": "Which shelf holds Parcel A?",
+              "criteria": {"a": "B2", "b": "C3"}})
+```
+
+The pip package only runs the model. Training code, data tools and evaluation scripts live in this repository.
+
+**Multiple choice**
+
+```json
+{"state": "Parcel A is on shelf B2. Parcel B is on shelf C3.",
+ "question": {"type": "choice", "instructions": "Which shelf holds Parcel A?",
+              "criteria": {"a": "B2", "b": "C3"}}}
+```
+```json
+{"confidence": 0.99996, "abstain": false, "choice": "a",
+ "probabilities": {"a": 0.99996, "b": 0.00004}}
+```
+
+**Yes/no**
+
+```json
+{"state": "Refunds need a receipt and a purchase within 30 days. The customer bought it 12 days ago and has no receipt.",
+ "question": {"type": "noul", "instructions": "Should the refund be issued?"}}
+```
+```json
+{"confidence": 0.933, "abstain": false, "noul": 0.067}
+```
+
+`noul` is the probability of yes.
+
+**Score**
+
+```json
+{"state": "The app crashes for every user when they open the payments page. No data is lost.",
+ "question": {"type": "score", "instructions": "How severe is this bug?",
+              "criteria": ["Cosmetic only", "Minor, with a workaround", "A core feature is blocked", "Data loss or harm"]}}
+```
+```json
+{"confidence": 0.431, "abstain": false, "score": 1.65,
+ "probabilities": [0.019, 0.431, 0.428, 0.122]}
+```
+
+`score` is the expected level (0 to 3 here). A low `confidence` like this one means the model is unsure; send those cases to a person.
+
+## Results
+
+| Evaluation | Accuracy |
+| --- | ---: |
+| Separate held-out test split (5,000) | 80.2% |
+| — multiple choice / yes-no / score | 80.8% / 88.0% / 65.8% |
+| Training dataset's test split (5,000) | 80.9% |
+| JevBench public (231) | 57.6% |
+| — easy / standard / hard | 97.9% / 66.7% / 34.2% |
+
+On the held-out split, yes/no AUC is 0.953 and expected calibration error is 0.037. At confidence 0.95 or higher the model answers 48.9% of questions and is right on 98.0% of them.
+
+## Set up for training
+
+```bash
+git clone https://github.com/kartikg-gk/tevro
+cd tevro
+uv sync
+uv run pytest -q tests
+```
+
+Python 3.12, managed with [uv](https://docs.astral.sh/uv/). On Windows, `uv sync` installs the CUDA 13.0 build of PyTorch; on Linux it uses the default PyPI build. The tests use a tiny randomly initialised model and run on CPU.
+
+## Train
+
+```bash
+uv run python train.py configs/smoke.yaml
+uv run python train.py configs/check_20k.yaml
+uv run python train.py configs/release_12k.yaml
+```
+
+| Config | What it does | Hardware |
+| --- | --- | --- |
+| `smoke.yaml` | Overfits 10 questions; train accuracy should reach 1.0 | any GPU |
+| `check_20k.yaml` | 20,000 questions, 1,000 test, bf16 weights, micro-batch 2 × 32 | 6 GB laptop GPU, ~45 min |
+| `release_12k.yaml` | The released model: 12,000 updates of 64 questions, fp32 | A100 80 GB, ~6.5 h |
+| `finetune_rules.yaml`, `finetune_v2.yaml` | Short fine-tunes of the released model on a 20K rule-reasoning mix (`--resume` the release checkpoint); ~10 min | A100 80 GB |
+
+Each run writes `runs/<name>-<timestamp>/` with the config, `metrics.jsonl`, `summary.json` and checkpoints (`best`, `latest`, `final`; LoRA adapter and head only). Every training log line shows `eta_min`, the projected minutes left. `--resume <run or checkpoint>` reloads the weights and starts a fresh optimizer.
+
+### On Modal
+
+```bash
+modal run --detach modal_app.py --config configs/release_12k.yaml
+modal run modal_app.py::evaluate --checkpoint runs/<run>/checkpoints/final --out eval.txt
+```
+
+`modal_app.py` uploads this folder, keeps `runs/` and the Hugging Face cache on Modal Volumes, runs on an A100 80 GB (set `TEVRO_GPU` to change it), and caps a training job at 9 hours.
+
+## Evaluate
+
+```bash
+uv run python inference.py runs/<run>/checkpoints/final --max_questions 5000
+uv run python jevbench_eval.py runs/<run>/checkpoints/final <jevbench>/datasets/public
+```
+
+The first scores a checkpoint on the training dataset's test split: accuracy per question type, yes/no AUC, ECE, Brier score, reliability and resolution, and accuracy at confidence 0.9 and 0.95. The second scores the 231 public [JevBench](https://github.com/fstandhartinger/jevbench) items by tier.
+
+## Data
+
+Each row is one state with one or more questions:
+
+```json
+{
+  "state": "Parcel A is on shelf B2. Parcel B is on shelf C3.",
+  "questions_json": "{\"q1\": {\"type\": \"choice\", \"instructions\": \"Which shelf holds Parcel A?\", \"criteria\": {\"a\": \"B2\", \"b\": \"C3\"}, \"label\": \"a\"}}"
+}
+```
+
+- `choice`: `criteria` maps option keys to descriptions; `label` is a key.
+- `noul` (yes/no): `label` is `true` or `false`.
+- `score`: `criteria` is an ordered list of levels; `label` is the level index.
+
+By default the code loads [kartikg1/tevro-decision-v1](https://huggingface.co/datasets/kartikg1/tevro-decision-v1) (250,000 training and 5,000 test questions; private). Point `data.question_files` at your own parquet files (`train` and `test`) to train on your data. `audit.py` reports label-position, length and negation shortcuts before you spend GPU time.
+
+## Model
+
+| Part | Setting |
+| --- | --- |
+| Base | Qwen3-0.6B, first 20 of 28 layers |
+| Adapter | LoRA rank 8, alpha 16, on q/k/v/o of layers 6–19 |
+| Input | task hint + question + state, then each option as its own span starting at the same position, then `The answer is:` |
+| Head | question-type embedding on every token, 2 bidirectional transformer layers (1024-d, 16 heads, dropout 0.1), scalar scorer at each option's last token |
+| Loss | cross-entropy over the options |
+| Optimizer | AdamW, LoRA lr 1e-4, head lr 1e-3, weight decay 0.01, cosine schedule, 5% warmup, gradient clip 1.0 |
+
+Options cannot attend to each other, so reordering them only reorders the probabilities. The test suite checks this.
+
+## Repository layout
+
+```
+tevro/              the pip package: copies of network.py and tokenization.py, plus predict.py and cli.py
+tokenization.py     input layout: positions, attention mask, read indices
+network.py          base model, LoRA, decision head
+dataset.py          loading, question expansion, batching
+train.py            training loop
+inference.py        evaluation metrics and single answers
+calibrate.py        temperature fitting, abstention thresholds, calibration metrics
+jevbench_eval.py    JevBench public evaluation
+audit.py            data shortcut checks
+synthetic.py        procedural decision questions
+data_variants.py    paired question variations
+reasoning_data.py   procedural rule-reasoning questions (policies, payouts, counting)
+build_mix.py        builds a 20K fine-tuning mix from generated, public and replayed questions
+llm_items.py        writes hard items with an LLM through OpenRouter and keeps those a blind second pass agrees with (needs `pydantic` and `OPENROUTER_API_KEY`)
+modal_app.py        Modal training and evaluation
+configs/            run settings
+tests/              layout, invariance, loss and data tests
+```
+
+## License
+
+MIT. The base model, Qwen3-0.6B, is released under its own license.

tevro-0.1.0/pyproject.toml

@@ -0,0 +1,60 @@
+[build-system]
+requires = ["hatchling>=1.27"]
+build-backend = "hatchling.build"
+
+[project]
+name = "tevro"
+version = "0.1.0"
+description = "A small decision model: picks one option and gives a probability for every option, in one forward pass"
+readme = "README.md"
+license = "MIT"
+license-files = ["LICENSE"]
+requires-python = ">=3.10"
+authors = [{ name = "Kartik Gadad" }]
+keywords = ["decision model", "classification", "lora", "qwen3", "calibration"]
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+dependencies = [
+    "huggingface_hub>=0.25.0",
+    "peft>=0.21.0",
+    "safetensors>=0.6.0",
+    "torch>=2.14.0",
+    "transformers>=5.17.0",
+]
+
+[project.urls]
+Homepage = "https://github.com/kartikg-gk/tevro"
+Model = "https://huggingface.co/kartikg1/tevro-decider"
+
+[project.scripts]
+tevro = "tevro.cli:main"
+
+[tool.hatch.build.targets.wheel]
+packages = ["tevro"]
+
+[tool.hatch.build.targets.sdist]
+include = ["tevro", "README.md", "LICENSE", "pyproject.toml"]
+
+# Training and data tools; not part of the published package. uv installs both groups by default.
+[dependency-groups]
+dev = [
+    "pytest>=8.0.0",
+]
+train = [
+    "datasets>=5.0.1",
+    "pyyaml>=6.0.3",
+]
+
+[tool.uv]
+default-groups = ["dev", "train"]
+
+# Windows has no CUDA torch on PyPI; take it from the PyTorch CUDA 13.0 index (Linux keeps PyPI's CUDA build)
+[tool.uv.sources]
+torch = [{ index = "pytorch-cu130", marker = "sys_platform == 'win32'" }]
+
+[[tool.uv.index]]
+name = "pytorch-cu130"
+url = "https://download.pytorch.org/whl/cu130"
+explicit = true

tevro-0.1.0/tevro/__init__.py

@@ -0,0 +1,13 @@
+"""Tevro: a small decision model. Give it a state, a question and options; it returns a probability per option."""
+
+__version__ = "0.1.0"
+__all__ = ["DEFAULT_MODEL", "Decider", "load"]
+
+
+def __getattr__(name):
+    # Lazy, so `import tevro` and `tevro --help` do not load the model stack or the tokenizer
+    if name in __all__:
+        from tevro import predict
+
+        return getattr(predict, name)
+    raise AttributeError(f"module 'tevro' has no attribute {name!r}")

tevro-0.1.0/tevro/cli.py

@@ -0,0 +1,24 @@
+"""Command line: `tevro predict request.json`."""
+import argparse
+import json
+import sys
+
+
+def main(argv=None):
+    parser = argparse.ArgumentParser(prog="tevro", description="Tevro decision model")
+    sub = parser.add_subparsers(dest="command", required=True)
+    p = sub.add_parser("predict", help="answer one question from a JSON file")
+    p.add_argument("input", help='JSON file with "state" and "question" (use - to read stdin)')
+    p.add_argument("--model", default=None, help="Hugging Face repo id or local folder (default: kartikg1/tevro-decider)")
+    p.add_argument("--device", default=None, help="cpu, cuda or mps (default: best available)")
+    args = parser.parse_args(argv)
+
+    from tevro.predict import DEFAULT_MODEL, load  # imported here so --help does not load the model stack
+
+    request = json.load(sys.stdin if args.input == "-" else open(args.input, encoding="utf-8"))
+    model = load(args.model or DEFAULT_MODEL, args.device)
+    print(json.dumps(model.answer(request["state"], request["question"]), indent=2))
+
+
+if __name__ == "__main__":
+    main()

tevro-0.1.0/tevro/inputs.py

@@ -0,0 +1,64 @@
+"""Turns question dicts into model inputs. Same logic the training code uses, without the dataset dependencies."""
+import torch
+
+from tevro.tokenization import tokenizer
+
+TASK_TYPES = {"choice": 0, "noul": 1, "score": 2}
+TASK_NAMES = {v: k for k, v in TASK_TYPES.items()}
+
+
+def question_to_choices(question):
+    """Turns any of the 3 question types into a list of option strings (+ label index if present)."""
+    if question["type"] == "choice":
+        keys = list(question["criteria"])
+        choices = [f"{key}: {question['criteria'][key]}" if question["criteria"][key] else key for key in keys]
+        label = keys.index(question["label"]) if "label" in question else None
+    elif question["type"] == "score":
+        # Choices share position ids, so the level index is written into the text to keep the ordering
+        choices = [f"{i}: {level}" for i, level in enumerate(question["criteria"])]
+        label = question.get("label")
+    else:  # noul -> P(Yes)
+        choices = ["No", "Yes"]
+        label = int(question["label"]) if "label" in question else None
+    return choices, label
+
+
+def collate_fn(batch):
+    B = len(batch)
+    L = max(len(ex["input_ids"]) for ex in batch)
+    C = max(len(ex["choice_read_idx"]) for ex in batch)
+
+    input_ids = torch.full((B, L), tokenizer.pad_token_id, dtype=torch.long)
+    position_ids = torch.zeros((B, L), dtype=torch.long)
+    # Padded query rows attend only to themselves, otherwise their softmax row is all -inf -> NaN
+    attention_mask = torch.eye(L).repeat(B, 1, 1)
+    token_mask = torch.zeros((B, L), dtype=torch.bool)
+    choice_read_idx = torch.zeros((B, C), dtype=torch.long)
+    choice_mask = torch.zeros((B, C), dtype=torch.bool)
+    answer_end_idx = torch.zeros(B, dtype=torch.long)
+
+    for b, ex in enumerate(batch):
+        n, c = len(ex["input_ids"]), len(ex["choice_read_idx"])
+        token_mask[b, :n] = True
+        input_ids[b, :n] = torch.tensor(ex["input_ids"])
+        position_ids[b, :n] = torch.tensor(ex["position_ids"])
+        attention_mask[b, :n, :n] = ex["attention_mask"]
+        choice_read_idx[b, :c] = torch.tensor(ex["choice_read_idx"])
+        choice_mask[b, :c] = True
+        answer_end_idx[b] = ex["answer_end_idx"]
+
+    attention_mask = torch.where(attention_mask == 0, -float("inf"), 0.0).unsqueeze(1)  # [B, 1, L, L]
+
+    out = {
+        "input_ids": input_ids,
+        "position_ids": position_ids,
+        "attention_mask": attention_mask,
+        "token_mask": token_mask,
+        "choice_read_idx": choice_read_idx,
+        "choice_mask": choice_mask,
+        "answer_end_idx": answer_end_idx,
+        "task_type": torch.tensor([ex["task_type"] for ex in batch], dtype=torch.long),
+    }
+    if all(ex["label"] is not None for ex in batch):
+        out["labels"] = torch.tensor([ex["label"] for ex in batch], dtype=torch.long)
+    return out

tevro-0.1.0/tevro/network.py

@@ -0,0 +1,137 @@
+import math
+
+import torch
+import torch.nn as nn
+from peft import LoraConfig, get_peft_model
+from transformers import Qwen3Model
+
+
+class ChoiceHead(nn.Module):
+    """Self-attention over [task token, choice vectors..., answer vector], then a bilinear score per choice.
+
+    No positional encoding: permuting the choices permutes the logits, so any number of choices works.
+    """
+
+    def __init__(self, hidden_dim, num_task_types, new_dim=512, num_layers=2, num_heads=8):
+        super().__init__()
+        self.new_dim = new_dim
+        self.task_embedding = nn.Embedding(num_task_types, new_dim)
+        self.input_proj = nn.Linear(hidden_dim, new_dim)
+        layer = nn.TransformerEncoderLayer(new_dim, num_heads, 4 * new_dim, dropout=0.0,
+                                           activation="gelu", batch_first=True, norm_first=True)
+        self.layers = nn.TransformerEncoder(layer, num_layers, enable_nested_tensor=False)
+        self.final_norm = nn.LayerNorm(new_dim)
+        self.answer_proj = nn.Linear(new_dim, new_dim)
+        self.choice_proj = nn.Linear(new_dim, new_dim)
+
+    def forward(self, choice_embeddings, answer_embedding, choice_mask, task_type):
+        # choice_embeddings: [B, C, H], answer_embedding: [B, H], choice_mask: [B, C], task_type: [B]
+        B = choice_embeddings.shape[0]
+        x = torch.cat([
+            self.task_embedding(task_type).unsqueeze(1),
+            self.input_proj(choice_embeddings),
+            self.input_proj(answer_embedding).unsqueeze(1),
+        ], dim=1)  # [B, C+2, D]
+
+        always_valid = torch.ones(B, 1, dtype=torch.bool, device=choice_mask.device)
+        ignore = ~torch.cat([always_valid, choice_mask, always_valid], dim=1)
+        x = self.final_norm(self.layers(x, src_key_padding_mask=ignore))
+
+        answer = self.answer_proj(x[:, -1])    # [B, D]
+        choices = self.choice_proj(x[:, 1:-1])  # [B, C, D]
+        logits = torch.einsum("bd,bcd->bc", answer, choices) / math.sqrt(self.new_dim)
+        return logits.masked_fill(~choice_mask, -float("inf"))
+
+    def score(self, hidden, token_mask, choice_read_idx, choice_mask, answer_end_idx, task_type):
+        rows = torch.arange(hidden.shape[0], device=hidden.device)
+        return self(hidden[rows[:, None], choice_read_idx], hidden[rows, answer_end_idx], choice_mask, task_type)
+
+
+class PostHead(nn.Module):
+    """Decision head over every token: type embedding added to all hidden states, a few bidirectional
+    transformer layers over the whole sequence, then a scalar scorer at each choice's read token.
+
+    No answer token is needed: the bidirectional layers let every choice see the state and the other
+    choices before scoring. No positional encoding, and the trunk gives choices shared positions,
+    so permuting the choices permutes the logits.
+    """
+
+    def __init__(self, hidden_dim, num_task_types, num_layers=2, dropout=0.1):
+        super().__init__()
+        num_heads = max(1, hidden_dim // 64)
+        self.task_embedding = nn.Embedding(num_task_types, hidden_dim)
+        layer = nn.TransformerEncoderLayer(hidden_dim, num_heads, 4 * hidden_dim, dropout,
+                                           batch_first=True, norm_first=True)
+        self.layers = nn.TransformerEncoder(layer, num_layers, enable_nested_tensor=False)
+        self.scorer = nn.Sequential(nn.LayerNorm(hidden_dim), nn.Linear(hidden_dim, hidden_dim), nn.GELU(),
+                                    nn.Linear(hidden_dim, 1))
+
+    def score(self, hidden, token_mask, choice_read_idx, choice_mask, answer_end_idx, task_type):
+        h = hidden + self.task_embedding(task_type)[:, None, :]
+        h = self.layers(h, src_key_padding_mask=~token_mask)
+        rows = torch.arange(h.shape[0], device=h.device)
+        logits = self.scorer(h[rows[:, None], choice_read_idx]).squeeze(-1)  # [B, C]
+        return logits.masked_fill(~choice_mask, -float("inf"))
+
+
+class TevroNetwork(nn.Module):
+    def __init__(self, backbone, head):
+        super().__init__()
+        self.backbone = backbone
+        self.head = head
+
+    def forward(self, input_ids, position_ids, attention_mask, token_mask, choice_read_idx, choice_mask,
+                answer_end_idx, task_type, **kwargs):
+        output = self.backbone(input_ids=input_ids, position_ids=position_ids, attention_mask=attention_mask,
+                               use_cache=False)
+        hidden = output.last_hidden_state.float()
+        # Head runs in fp32 even under bf16 autocast: bf16 logits are too coarse for a stable loss
+        with torch.autocast(device_type=hidden.device.type, enabled=False):
+            return self.head.score(hidden, token_mask, choice_read_idx, choice_mask, answer_end_idx, task_type)
+
+
+def load_backbone(model_name, num_layers=None, dtype="float32"):
+    # Only the first num_layers decoder layers are built and loaded (null = all of them).
+    # bfloat16 halves the frozen weights for small GPUs; compute already runs in bf16 under autocast
+    overrides = {} if num_layers is None else {"num_hidden_layers": num_layers}
+    return Qwen3Model.from_pretrained(model_name, dtype=getattr(torch, dtype), **overrides)
+
+
+def lora_target_regex(num_layers, last_k_layers):
+    layer_ids = "|".join(str(i) for i in range(num_layers - last_k_layers, num_layers))
+    return rf".*layers\.({layer_ids})\.self_attn\.(q_proj|k_proj|v_proj|o_proj)"
+
+
+def add_lora(backbone, r, alpha, dropout, last_k_layers):
+    if last_k_layers == 0:  # frozen Qwen, only the head trains
+        backbone.requires_grad_(False)
+        return backbone
+    config = LoraConfig(
+        r=r,
+        lora_alpha=alpha,
+        lora_dropout=dropout,
+        target_modules=lora_target_regex(backbone.config.num_hidden_layers, last_k_layers),
+    )
+    return get_peft_model(backbone, config)
+
+
+def build_network(model_name, num_layers, lora_config, head_config, num_task_types, backbone=None,
+                  dtype="float32"):
+    if backbone is None:
+        backbone = load_backbone(model_name, num_layers, dtype)
+    backbone = add_lora(backbone, lora_config["r"], lora_config["alpha"], lora_config["dropout"],
+                        lora_config["last_k_layers"])
+    head = build_head(head_config, backbone.config.hidden_size, num_task_types)
+    return TevroNetwork(backbone, head)
+
+
+HEADS = {"choicehead": ChoiceHead, "posthead": PostHead}
+
+
+def build_head(head_config, hidden_dim, num_task_types):
+    """head.type: choicehead = attention over [task, choices, answer token] + bilinear score |
+    posthead = bidirectional layers over every token + scalar scorer per choice."""
+    head_config = dict(head_config)
+    head_type = head_config.pop("type", "choicehead")
+    head_config.pop("num_task_types", None)
+    return HEADS[head_type](hidden_dim, num_task_types, **head_config)

tevro-0.1.0/tevro/predict.py

@@ -0,0 +1,85 @@
+"""Load a trained Tevro model and answer one question at a time."""
+import json
+from pathlib import Path
+
+import torch
+from huggingface_hub import snapshot_download
+from peft import PeftModel
+from safetensors.torch import load_file
+
+from tevro.inputs import TASK_TYPES, collate_fn, question_to_choices
+from tevro.network import TevroNetwork, build_head, load_backbone
+from tevro.tokenization import encode_example
+
+DEFAULT_MODEL = "kartikg1/tevro-decider"
+
+
+def get_device():
+    if torch.cuda.is_available():
+        return torch.device("cuda")
+    if torch.backends.mps.is_available():
+        return torch.device("mps")
+    return torch.device("cpu")
+
+
+def resolve(model):
+    """A local folder is used as is; anything else is treated as a Hugging Face repo id."""
+    if Path(model).is_dir():
+        return Path(model)
+    return Path(snapshot_download(model, allow_patterns=["meta.json", "head.safetensors", "lora/*"]))
+
+
+class Decider:
+    """A loaded model. `answer(state, question)` returns the chosen option and a probability for every option."""
+
+    def __init__(self, network, meta, device):
+        self.network = network.to(device).eval()
+        self.meta = meta
+        self.device = torch.device(device)
+
+    @torch.no_grad()
+    def answer(self, state, question):
+        """state: text or any JSON-able object. question: {"type": "choice"|"noul"|"score", "instructions": ...,
+        "criteria": {key: description} for choice, a list of ordered levels for score, optional for noul}."""
+        if not isinstance(state, str):
+            state = json.dumps(state, ensure_ascii=False)
+        choices, _ = question_to_choices(question)
+        example = encode_example(state, question["instructions"], choices, question["type"],
+                                 self.meta["max_state_tokens"], self.meta["max_choice_tokens"],
+                                 self.meta.get("use_answer_prompt", True))
+        example.update(task_type=TASK_TYPES[question["type"]], label=None)
+        batch = {k: v.to(self.device) for k, v in collate_fn([example]).items()}
+
+        autocast = torch.autocast(self.device.type, dtype=torch.bfloat16, enabled=self.device.type in ("cuda", "mps"))
+        with autocast:
+            logits = self.network(**batch)[0].float().cpu()
+        calibration = self.meta.get("calibration") or {}
+        logits = logits / calibration.get("temperatures", {}).get(question["type"], 1.0)
+        probs = torch.softmax(logits, -1)
+        confidence = probs.max().item()
+        threshold = calibration.get("thresholds", {}).get(question["type"], 0.0)
+        out = {"confidence": confidence, "abstain": confidence < threshold}
+        probs = probs.tolist()
+
+        if question["type"] == "choice":
+            keys = list(question["criteria"])
+            out.update(choice=keys[max(range(len(keys)), key=probs.__getitem__)], probabilities=dict(zip(keys, probs)))
+        elif question["type"] == "score":
+            out.update(score=sum(i * p for i, p in enumerate(probs)), probabilities=probs)
+        else:
+            out.update(noul=probs[1])
+        return out
+
+
+def load(model=DEFAULT_MODEL, device=None):
+    """Loads the base model, the LoRA adapter and the decision head. `model` is a Hugging Face repo id or a folder."""
+    folder = resolve(model)
+    meta = json.loads((folder / "meta.json").read_text())
+    device = torch.device(device) if device else get_device()
+
+    backbone = load_backbone(meta["model_name"], meta.get("num_layers"), meta.get("dtype", "float32"))
+    if (folder / "lora").exists():
+        backbone = PeftModel.from_pretrained(backbone, folder / "lora")
+    head = build_head(meta["head_config"], backbone.config.hidden_size, meta["head_config"]["num_task_types"])
+    head.load_state_dict(load_file(str(folder / "head.safetensors"), device="cpu"))
+    return Decider(TevroNetwork(backbone, head), meta, device)

tevro-0.1.0/tevro/tokenization.py

@@ -0,0 +1,87 @@
+import os
+
+import torch
+from transformers import AutoTokenizer
+
+os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")  # the tokenizer is shared with DataLoader worker processes
+
+model_name = "Qwen/Qwen3-0.6B"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+OPTION_OPEN = tokenizer.encode("<option>")
+OPTION_CLOSE = tokenizer.encode("</option>")
+ANSWER_PROMPT = tokenizer.encode("The answer is:")
+
+# Short natural-language task hint
+TASK_PROMPTS = {
+    "choice": "Choose the best option.",
+    "score": "Rate it on the given scale.",
+    "noul": "Answer yes or no.",
+}
+
+
+def truncate(text, max_tokens):
+    tokens = tokenizer.encode(text)
+    if len(tokens) <= max_tokens:
+        return text
+    return tokenizer.decode(tokens[:max_tokens])
+
+
+def build_question(state, instructions, task_type, max_state_tokens):
+    # Task hint and question come before the state, so every state token is read knowing what is asked
+    return f"{TASK_PROMPTS[task_type]}\nQuestion: {instructions}\n\n{truncate(state, max_state_tokens)}\n"
+
+
+def build_layout(question_ids, choice_ids, use_answer_prompt=True):
+    """One sequence: question, then every choice wrapped in <option>...</option>, then "The answer is:".
+
+    Returns token ids, position ids, a [L, L] attention mask (1 = may attend) and the read indices.
+    Every choice starts at the same position and cannot see the other choices, so the choices are
+    read in parallel: reordering them only reorders the outputs. The answer tokens come after the
+    longest choice and see everything, which is where the choices are compared.
+    """
+    q = len(question_ids)
+    input_ids = list(question_ids)
+    position_ids = list(range(q))
+    spans = []
+    for ids in choice_ids:
+        wrapped = OPTION_OPEN + ids + OPTION_CLOSE
+        start = len(input_ids)
+        input_ids += wrapped
+        position_ids += range(q, q + len(wrapped))
+        spans.append((start, len(input_ids)))  # [start, end)
+
+    answer_start = len(input_ids)
+    answer_pos = max(position_ids) + 1
+    answer_ids = ANSWER_PROMPT if use_answer_prompt else []
+    input_ids += answer_ids
+    position_ids += range(answer_pos, answer_pos + len(answer_ids))
+    L = len(input_ids)
+
+    mask = torch.zeros((L, L))
+    mask[:q, :q] = torch.tril(torch.ones(q, q))           # question: causal
+    for start, end in spans:
+        n = end - start
+        mask[start:end, :q] = 1                            # each choice reads the question and state
+        mask[start:end, start:end] = torch.tril(torch.ones(n, n))  # and itself, causally
+    n = L - answer_start
+    mask[answer_start:, :answer_start] = 1                 # answer reads everything before it
+    mask[answer_start:, answer_start:] = torch.tril(torch.ones(n, n))
+
+    # Choices are read at their last content token, not at the closing tag
+    choice_read_idx = [end - len(OPTION_CLOSE) - 1 for _, end in spans]
+    return {
+        "input_ids": input_ids,
+        "position_ids": position_ids,
+        "attention_mask": mask,
+        "choice_read_idx": choice_read_idx,
+        "answer_end_idx": L - 1,
+    }
+
+
+def encode_example(state, instructions, choices, task_type, max_state_tokens, max_choice_tokens,
+                   use_answer_prompt=True):
+    question = build_question(state, instructions, task_type, max_state_tokens)
+    choices = [truncate(choice, max_choice_tokens) for choice in choices]
+    return build_layout(tokenizer.encode(question), [tokenizer.encode(choice) for choice in choices],
+                        use_answer_prompt)