parsecai/tod

Code: https://github.com/daseinlabs/tod · Base model is not included; accept the Gemma 4 licence on Hugging Face and it is pulled on first load.

Universal two-stage picker (JevBench)

Base model: google/gemma-4-12B-it (revision default), Apache-2.0. Two-stage picker: pointwise recall over unbounded N -> letter re-rank of the top-16 shortlist (research/28 §5-§9).

Configuration

• stage-1 kind: union (k=16, k_each=8)
• stage-1 adapter: none (frozen); stage-2 adapter: stage2_adapter; retriever: retriever
• attention: chunked_eager; dtype: eval float32 / serve bfloat16; max_tokens: 49152
• prompt versions: stage-1 pointwise.v1, stage-2 cygnet.v1; option style label_desc
• temperatures: T1=1.7237, T1_ret=0.05, T2=1.4778

Results

Public JevBench (231, text-only)

variant	acc	easy / orig / hard	Brier	ECE-db
release (trained stage 2)	0.844	1.000 / 0.917 / 0.730	0.194	0.0328
frozen floor	0.861	1.000 / 0.944 / 0.748	0.195	0.0437

High-N (per set, calibrated two-stage k=16)

set	N	acc	recall@16
banking77	799	0.735	n/a
clinc150	n/a	n/a	n/a
lexglue	n/a	n/a	n/a
tasksource	n/a	n/a	n/a
combined		0.735	0.963

Vision (Δacc = with-image − without-image, item bootstrap)

set	acc with	acc without	Δacc	earns keep
scienceqa	0.90	0.73	+0.17	True
ai2d	0.87	0.62	+0.25	True
chartqa	0.97	0.79	+0.18	True

Usage

from tod.release.load import load_picker

picker = load_picker("path/to/this/dir")          # reads picker.json, rebuilds the TwoStageScorer

probs = picker.predict(                            # one calibrated probability per option
    state="Customer: my card was charged twice.",
    options={"refund": "wants a refund", "dispute": "disputes a charge"},  # or a plain label list
    question="Pick the user's intent.",            # str, full question dict, or None
)
# probs is aligned to the options order and sums to 1, e.g. [0.71, 0.29]

predict(state, options, question=None, images=None) runs stage 1 (pointwise recall over all N; unioned with the retriever when present) then the stage-2 letter re-rank over the top-k shortlist, applies the calibrated temperatures from picker.json, and returns the probability for EVERY option in the given order (twostage_final_probs). N is never capped; for N ≤ k the picker is exactly the pure letter read. Pass images (paths) to score a multimodal state through the native processor. Serve dtype is bf16; the exact-eval dtype is fp32 (see picker.json).

Licence & data

• Model weights: Apache-2.0 (see LICENSE). Base weights are Google Gemma 4 (Apache-2.0); this release modifies them via LoRA adapters. Attribution and the "modified files" statement are in NOTICE.
• Gemma Prohibited Use Policy: https://ai.google.dev/gemma/prohibited_use_policy
• Training-data attribution: see data_pipeline/mixes/ATTRIBUTIONS.md in the source repository.