Publish mini-Jev step 010626 interim checkpoint

LICENSE

@@ -1,3 +1,4 @@
+
                                  Apache License
                            Version 2.0, January 2004
                         http://www.apache.org/licenses/
@@ -174,3 +175,28 @@
       of your accepting any such warranty or additional liability.
 
    END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

README.md

@@ -1,277 +1,75 @@
 ---
 license: apache-2.0
-base_model:
-  - Qwen/Qwen3-0.6B
-library_name: pytorch
+base_model: Qwen/Qwen3-0.6B
 tags:
-  - jev
   - mini-jev
   - decision-model
-  - agent
-  - agentic-ai
   - tool-selection
-  - tool-routing
-  - model-routing
-  - function-calling
-  - agent-control
-  - local-inference
-  - open-weights
-  - qwen
+  - tool-use
+  - agents
+  - action-selection
+  - decision-making
+  - finite-choice
+  - qwen3
+  - peft
 ---
 
-# mini-Jev
+# mini-Jev Qwen3-0.6B - step 10,626
 
-**A Qwen3-0.6B-based decision model for tool selection, routing, and agent-control experiments.**
+mini-Jev is a **0.6B agentic decision model specialized for finite-choice tool and action selection**. Give it a state, a question, and candidate actions; it returns a probability for each candidate. It uses a Qwen3-0.6B backbone with an adapter and a compact decision head. Candidate scoring is permutation equivariant, and the model returns finite-choice probabilities rather than generated text or tool arguments.
 
-mini-Jev is an experimental decision model designed to make structured control decisions inside an agent loop:
+This is an **early 10,626-update checkpoint**. Training is planned through 50,000 updates; more training and broader decision supervision are coming. This release candidate is not the final checkpoint.
 
-```text
-state + available candidates → probabilities + selected decision
-```
+## Measured results
 
-The current release:
-- Uses [Qwen3-0.6B](https://huggingface.co/Qwen/Qwen3-0.6B) as a backbone.
-- Keeps the Qwen backbone frozen.
-- Adds a lightweight trained decision head (~1.1 MB) to score candidate actions directly.
-- Performs candidate scoring and selection; it is not designed for text generation.
-- Serves as an experimental baseline and demo for fast, structured agent decision-making.
+| Evaluation | Result | Scope |
+|---|---:|---|
+| BFCL V1 Multiple Function - Jev function-selection adaptation | **96.5% (193/200)** | Selects the supplied function identity from 2-4 definitions. **Not an official BFCL leaderboard score**; does not generate arguments or execute calls. |
+| Strict held-out Choice | **65.6% (63/96)** | A selected, bucket-balanced validation slice, not the full validation split. |
+| Bespoke OOD typed decisions | **72.0% (36/50)** | Small synthetic sanity check with straightforward distractors; not a public benchmark. |
+| LocalLLaMA/typed-decisions, full test set | **34.0% (680/2,000)** | Choice 27.0%, Noul 52.2%, Score 25.6%. General typed decisions remain a weakness. |
 
-> **Note on naming:** The model implementation was originally developed under the internal name **ODM Mini v1**. That name remains in class names (`ODMMiniModel`), loader files (`odm_mini.py`), and `config.json` for backwards compatibility.
+![BFCL selection accuracy at step 0 and step 10,626](assets/bfcl-step-comparison.svg)
 
----
-
-## 1. What is mini-Jev?
-
-AI agents frequently encounter situations where they must make small, structured decisions rather than generate free-form text. Common examples include:
-
-- **Tool selection:** Deciding which tool or external integration to invoke.
-- **Function routing:** Selecting the specific API or method to execute next.
-- **Next-action prediction:** Choosing between continuing investigation, asking for clarification, or concluding.
-- **Continue / finish decisions:** Deciding whether an assigned task is satisfied or requires another step.
-- **Model routing:** Directing incoming tasks to appropriate candidate downstream models.
-
-Rather than invoking a large generative language model for every routing and control decision, mini-Jev explores using a small, specialized decision model to evaluate candidate options directly over agent state.
-
-*mini-Jev is an experimental baseline and is not intended for unmonitored production use.*
-
----
-
-## 2. Current model
-
-**Current release: Qwen3-0.6B-based mini-Jev v1**
-
-### Architecture
-
-```text
-Frozen Qwen3-0.6B → candidate representation → lightweight decision head → grouped softmax
-```
-
-1. **Backbone:** The base model is [`Qwen/Qwen3-0.6B`](https://huggingface.co/Qwen/Qwen3-0.6B), which remains completely frozen. The weights in this repository contain only the trained decision head; the Qwen backbone is downloaded separately by the loader.
-2. **Representation pooling:** State text and candidate descriptions are passed through the model. Only candidate-description tokens are mean-pooled. Candidate IDs, prompt formatting headers, state tokens, and padding tokens are excluded from the pooling mask.
-3. **Decision head:** A lightweight two-layer MLP:
-   - `Linear(1024, 256)`
-   - `GELU`
-   - `Linear(256, 1)`
-   - Total head parameters: **262,657** (~1.1 MB).
-4. **Scoring:** The resulting scalar logits are normalized across candidate options using a grouped softmax to yield probabilities, a selected candidate, confidence, and decision margin.
-
----
-
-## 3. Quick start
+![BFCL selection accuracy by candidate count](assets/bfcl-by-candidates.svg)
 
-### Installation
+![Accuracy snapshot across separate evaluations](assets/evaluation-summary.svg)
 
-```bash
-pip install torch transformers safetensors huggingface_hub
-```
+These evaluations have different datasets and denominators; their percentages are not directly comparable. The BFCL adaptation uses the [BFCL V1 Multiple Function AST examples](https://github.com/EnlightenedAI/BFCL/blob/main/berkeley-function-call-leaderboard/bfcl_eval/data/README.md). The typed-decision result uses the [LocalLLaMA/typed-decisions test split](https://huggingface.co/datasets/LocalLLaMA/typed-decisions).
 
-Download the loader module:
+## Intended use and limits
 
-```bash
-hf download samatv256/mini-Jev odm_mini.py --local-dir .
-```
+The current strength is **finite-choice tool and action selection** from state and supplied candidates. The model is less reliable for general typed decisions and for large candidate sets. It is not a general probability or classification model, and its probabilities have not been calibrated for arbitrary domains. A selected 17+ candidate held-out slice scored 4/17, so large sets need particular care.
 
-### Basic usage
+## Use
 
-```python
-from odm_mini import ODMMiniModel
+A CUDA GPU with BF16 support is required for this 4-bit package. Install the tested dependencies from requirements.txt. The [Qwen3-0.6B base model](https://huggingface.co/Qwen/Qwen3-0.6B) is downloaded separately on first load; this repository contains only the adapter and decision head.
 
-model = ODMMiniModel.from_pretrained("samatv256/mini-Jev")
+    from inference import MiniJev
 
-choice = model.predict_choice(
-    state={"user_request": "Find the weather in Boston."},
-    candidates=[
-        {
-            "id": "weather.lookup",
-            "description": "Look up the current weather.",
-        },
-        {
-            "id": "calendar.list",
-            "description": "List calendar events.",
-        },
-    ],
-)
-
-print("Selected:", choice.selected)
-print("Probabilities:", choice.probabilities)
-```
-
-### Full snapshot download and usage
-
-You can also download all repository files to a local directory before loading:
-
-```python
-import sys
-from pathlib import Path
-from huggingface_hub import snapshot_download
-
-release_dir = Path(
-    snapshot_download(
-        repo_id="samatv256/mini-Jev",
-        allow_patterns=[
-            "README.md",
-            "model.safetensors",
-            "config.json",
-            "odm_mini.py",
-            "LICENSE",
+    model = MiniJev.load(".")
+    result = model.predict(
+        state={"user_goal": "Find tomorrow's weather forecast for Paris."},
+        question="Given the current state and available options,\nwhich option should be selected?",
+        question_type="choice",
+        answer_options=[
+            {"id": "weather", "type": "tool", "label": "get_weather_forecast",
+             "description": "Look up the weather forecast for a place and date."},
+            {"id": "invoice", "type": "tool", "label": "calculate_invoice_total",
+             "description": "Add amounts on an invoice."},
+            {"id": "email", "type": "tool", "label": "send_email",
+             "description": "Send an email message."},
         ],
     )
-)
-sys.path.insert(0, str(release_dir))
-
-from odm_mini import ODMMiniModel
-
-model = ODMMiniModel.from_pretrained(release_dir)
-
-state = {
-    "user_request": "Find the current weather in Boston.",
-    "available_context": "The user has not provided weather data.",
-}
-candidates = [
-    {
-        "id": "weather.lookup",
-        "description": "Look up current weather for a specified city.",
-    },
-    {
-        "id": "calendar.list",
-        "description": "List upcoming calendar events for the user.",
-    },
-    {
-        "id": "control.finish",
-        "description": "Finish because the request has already been completed.",
-    },
-]
-
-choice = model.predict_choice(state=state, candidates=candidates)
-print("Selected candidate:", choice.selected)
-print("Probabilities:", choice.probabilities)
-print("Confidence:", choice.confidence)
-print("Decision margin:", choice.decision_margin)
-print("Latency (ms):", choice.latency_ms)
-```
-
-`ODMMiniModel.from_pretrained()` accepts either the Hugging Face repo ID or a local directory path. CUDA runs in BF16 by default; CPU runs in FP32.
-
----
-
-## 4. Example use cases
-
-mini-Jev can be used for several structured agent-control patterns:
-
-- **Tool selection:** Choosing the right tool from a list of available integrations.
-- **Function routing:** Determining which endpoint or handler to call.
-- **Next-action prediction:** Selecting intermediate steps in a reasoning or execution plan.
-- **Agent-control research:** Investigating small-footprint models for local agent decision-making.
-- **Model-routing experiments:** Evaluating routing policies between different candidate models. *(Note: Model routing is a conceptual application of the candidate-selection interface, not a validated production routing capability.)*
-
-### Conceptual example
+    print(result["selected_id"])
+    print(result["options"])
 
-```text
-State:
-User wants to find the weather in Boston.
+Option IDs are bookkeeping and are excluded from model text. The model scores each supplied option and normalizes probabilities over the finite set. Input branches are limited to 8,192 tokens.
 
-Candidates:
-- web_search
-- weather_tool
-- calculator
-- finish
-
-Decision:
-weather_tool
-```
-
----
-
-## 5. Jev Decisions v1
-
-Researchers interested in training and evaluating decision models can explore the [Jev Decisions v1 dataset](https://huggingface.co/datasets/samatv256/jev-decisions-v1).
-
-- **Dataset card:** [https://huggingface.co/datasets/samatv256/jev-decisions-v1](https://huggingface.co/datasets/samatv256/jev-decisions-v1)
-- **Description:** A public dataset containing nearly 12 million canonical agent-decision records (over 6.2M choice-eligible decisions across 1.06M task trajectories) with explicit candidate sets, chosen actions, and eligibility flags.
-
-> **Important:** The current mini-Jev Qwen3-0.6B baseline was **NOT trained on Jev Decisions v1**. It is provided as an open project resource for community research.
-
----
-
-## 6. Current verified results
-
-The results below reflect the public Qwen3-0.6B baseline checkpoint:
-
-### Training configuration
-
-- **Training data:** 50,000 synthetic decision examples
-- **Backbone:** Frozen `Qwen/Qwen3-0.6B`
-- **Trained components:** DecisionHead weights only (grouped cross-entropy objective)
-- **Checkpoint:** Seed 41, epoch 4
-- **Temperature:** Fixed at 1.0; maximum softmax output is reported as confidence.
-
-### Synthetic held-out evaluation
-
-| Metric | Result |
-|---|---:|
-| Semantic Choice accuracy | **72.97%** |
-| Stress Choice accuracy | **67.64%** |
-| Counterfactual pair consistency | **67.12%** |
-
-- **Latency measurements:** On an NVIDIA GH200 using BF16 backbone inference and a shared-prefix KV-cache path, total latency was approximately **76–85 ms** for contexts of 256–1,024 state tokens across 3–16 candidates (measured range: 76.11–82.36 ms). Candidate representation extraction reached 831.3 candidates/second at batch size 512.
-
-### Real shadow-agent evaluation
-
-In an offline evaluation on real agent executions across **75 multi-step trajectories and 243 decisions**, the baseline achieved:
-
-- **Action accuracy:** **27.98%**
-- **Controller agreement:** **26.34%**
-
-These metrics are reported to transparently show the significant transfer gap between synthetic single-step decisions and dynamic, multi-step agent trajectories.
-
----
-
-## 7. Limitations
-
-> **mini-Jev v1 is an experimental research prototype and is not ready for unmonitored production agent control.**
-
-- **Limited real-agent transfer:** While performance on synthetic single-step benchmarks is moderate (~73%), accuracy drops substantially (~28%) in real multi-step agent environments.
-- **Premature completion:** The model exhibits a known failure mode of high-confidence premature `finish` decisions. After making partial progress, it frequently over-indexes on successful intermediate receipts and chooses `control.finish` before completing remaining steps.
-- **Uncalibrated confidence:** Grouped softmax probabilities indicate relative preference among the provided candidates, not calibrated real-world uncertainty.
-- **No autonomous control:** Do not use this baseline as the sole decision-maker for consequential actions or mission-critical workflows without a human or supervising controller in the loop.
-
----
-
-## 8. Public project links
-
-- **Dataset:** [Jev Decisions v1](https://huggingface.co/datasets/samatv256/jev-decisions-v1)
-- **Collection:** [mini-Jev Hugging Face Collection](https://huggingface.co/collections/samatv256/mini-jev-small-decision-models-for-ai-agents-6ab3e559bed2124120973013)
-
-Future work will continue exploring improved decision models and training on larger real agent/tool-use datasets.
-
----
+## Data attribution
 
-## Released files and license
+Decision supervision for this checkpoint draws in part on [Jev Decisions v1](https://huggingface.co/datasets/samatv256/jev-decisions-v1), a public dataset released under CC BY 4.0. Its public upstream datasets carry their own attribution and licensing terms; see the [dataset card](https://huggingface.co/datasets/samatv256/jev-decisions-v1) and [upstream license notes](https://huggingface.co/datasets/samatv256/jev-decisions-v1/blob/main/SOURCE_LICENSES.md).
 
-- `model.safetensors` — Trained DecisionHead weights only (262,657 parameters, ~1.1 MB)
-- `config.json` — Architecture specification and backbone reference
-- `odm_mini.py` — Inference loader module
-- `README.md` — Model card and usage documentation
-- `LICENSE` — Apache License 2.0
+## License
 
-The mini-Jev decision head and loader code are released under the Apache-2.0 License. The separately downloaded Qwen3-0.6B backbone is governed by its own license terms from the Qwen team.
+Apache-2.0. The Qwen3-0.6B base model is separately licensed by Qwen under Apache-2.0. See LICENSE and the [base model card](https://huggingface.co/Qwen/Qwen3-0.6B).

adapter/adapter_config.json

@@ -0,0 +1,79 @@
+{
+  "alora_invocation_tokens": null,
+  "alpha_pattern": {},
+  "arrow_config": null,
+  "auto_mapping": {
+    "base_model_class": "Qwen3Model",
+    "parent_library": "transformers.models.qwen3.modeling_qwen3"
+  },
+  "base_model_name_or_path": "Qwen/Qwen3-0.6B",
+  "bias": "none",
+  "corda_config": null,
+  "ensure_weight_tying": false,
+  "eva_config": null,
+  "exclude_modules": null,
+  "fan_in_fan_out": false,
+  "inference_mode": true,
+  "init_lora_weights": true,
+  "kasa_config": null,
+  "layer_replication": null,
+  "layers_pattern": null,
+  "layers_to_transform": null,
+  "loftq_config": {},
+  "lora_alpha": 32,
+  "lora_bias": false,
+  "lora_dropout": 0.0,
+  "lora_ga_config": null,
+  "megatron_config": null,
+  "megatron_core": "megatron.core",
+  "modules_to_save": null,
+  "monteclora_config": null,
+  "peft_type": "LORA",
+  "peft_version": "0.21.0",
+  "qalora_group_size": 16,
+  "r": 16,
+  "rank_pattern": {},
+  "revision": null,
+  "target_modules": [
+    "layers.22.self_attn.v_proj",
+    "layers.21.self_attn.v_proj",
+    "layers.27.self_attn.v_proj",
+    "layers.22.self_attn.k_proj",
+    "layers.27.self_attn.o_proj",
+    "layers.20.self_attn.v_proj",
+    "layers.21.self_attn.o_proj",
+    "layers.21.self_attn.q_proj",
+    "layers.24.self_attn.v_proj",
+    "layers.25.self_attn.o_proj",
+    "layers.26.self_attn.v_proj",
+    "layers.25.self_attn.q_proj",
+    "layers.27.self_attn.q_proj",
+    "layers.24.self_attn.q_proj",
+    "layers.27.self_attn.k_proj",
+    "layers.26.self_attn.q_proj",
+    "layers.26.self_attn.k_proj",
+    "layers.23.self_attn.q_proj",
+    "layers.24.self_attn.o_proj",
+    "layers.25.self_attn.k_proj",
+    "layers.23.self_attn.v_proj",
+    "layers.22.self_attn.o_proj",
+    "layers.20.self_attn.o_proj",
+    "layers.20.self_attn.k_proj",
+    "layers.23.self_attn.o_proj",
+    "layers.26.self_attn.o_proj",
+    "layers.25.self_attn.v_proj",
+    "layers.21.self_attn.k_proj",
+    "layers.22.self_attn.q_proj",
+    "layers.20.self_attn.q_proj",
+    "layers.23.self_attn.k_proj",
+    "layers.24.self_attn.k_proj"
+  ],
+  "target_parameters": null,
+  "task_type": null,
+  "trainable_token_indices": null,
+  "use_bdlora": null,
+  "use_dora": false,
+  "use_qalora": false,
+  "use_rslora": false,
+  "velora_config": null
+}

adapter/adapter_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:765e987436517f18ef902641a197c719a415292d082615ca891e38590fbffeda
+size 5251016

config.json

@@ -1,11 +1,16 @@
 {
-  "architectures": ["ODMMiniModel"],
-  "model_type": "odm_mini",
   "base_model": "Qwen/Qwen3-0.6B",
-  "hidden_size": 1024,
-  "head_inner_size": 256,
-  "pooling": "candidate_description_mean",
-  "temperature": 1.0,
-  "max_state_length": 4096,
-  "max_suffix_length": 256
+  "base_revision": "c1899de289a04d12100db370d81485cdf75e47ca",
+  "max_tokens": 8192,
+  "pooler": "mean",
+  "projection_dim": 256,
+  "set_layers": 2,
+  "set_heads": 4,
+  "quantization": {
+    "format": "nf4",
+    "double_quantization": true,
+    "compute_dtype": "bfloat16"
+  },
+  "adapter_dir": "adapter",
+  "head_file": "decision_head.safetensors"
 }

decision_head.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f06d990297b23ffabc7a0d5eba3670830c337b742f8e3a30fe3907783e0bac88
+size 7631364

inference.py

@@ -0,0 +1,262 @@
+"""Public, inference-only loader for the mini-Jev step 10,626 adapter and decision head."""
+from __future__ import annotations
+
+import json
+from collections.abc import Mapping, Sequence
+from pathlib import Path
+from typing import Any
+
+import torch
+from peft import PeftModel, prepare_model_for_kbit_training
+from safetensors.torch import load_file
+from torch import nn
+from transformers import AutoModel, AutoTokenizer, BitsAndBytesConfig
+
+
+def _text(value: Any) -> str:
+    if value is None:
+        return ""
+    if isinstance(value, str):
+        try:
+            value = json.loads(value)
+        except (ValueError, TypeError):
+            return value.strip()
+    if isinstance(value, (dict, list)):
+        return json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":"))
+    return str(value)
+
+
+def _section(name: str, content: str) -> str:
+    return f"<{name}>\n{content}\n</{name}>\n"
+
+
+def _option_text(option: Mapping[str, Any]) -> str:
+    label = option.get("label") or option.get("name")
+    if not label:
+        raise ValueError("each answer option needs a nonempty label")
+    fields = (
+        ("type", option.get("type") or option.get("option_type") or "action"),
+        ("label", label),
+        ("description", option.get("description")),
+        ("schema", option.get("schema") or option.get("parameters_json") or option.get("parameters")),
+    )
+    lines = [f"{name}: {_text(value)}" for name, value in fields if value is not None and _text(value)]
+    return _section("ANSWER_OPTION", "\n".join(lines))
+
+
+def _encode(tokenizer, value: str) -> list[int]:
+    return list(tokenizer.encode(value, add_special_tokens=False))
+
+
+def encode_branch(
+    tokenizer,
+    state: Mapping[str, Any] | str,
+    question: str,
+    question_type: str,
+    option: Mapping[str, Any],
+    max_tokens: int = 8192,
+) -> tuple[list[int], list[bool]]:
+    """Serialize one option branch; preserve the question and option within the 8K budget."""
+    if not 1 <= max_tokens <= 8192:
+        raise ValueError("max_tokens must be 1..8192")
+    if not question or not question_type:
+        raise ValueError("question and question_type must be nonempty")
+    if isinstance(state, str):
+        state = {"summary": state}
+    if not isinstance(state, Mapping):
+        raise TypeError("state must be a string or a mapping")
+    history = state.get("history") or []
+    if not isinstance(history, Sequence) or isinstance(history, (str, bytes)):
+        raise TypeError("state.history must be a sequence")
+    state_open = _encode(tokenizer, "<STATE>\n")
+    system = _encode(tokenizer, _section("SYSTEM", _text(state.get("system"))))
+    goal = _encode(tokenizer, _section("USER_GOAL", _text(state.get("user_goal"))))
+    summary_open = _encode(tokenizer, "<SUMMARY>\n")
+    summary_content = _encode(tokenizer, _text(state.get("summary")))
+    summary_close = _encode(tokenizer, "\n</SUMMARY>\n")
+    environment_open = _encode(tokenizer, "<ENVIRONMENT>\n")
+    environment_content = _encode(tokenizer, _text(state.get("environment")))
+    environment_close = _encode(tokenizer, "\n</ENVIRONMENT>\n")
+    open_history = _encode(tokenizer, "<HISTORY>\n")
+    close_history = _encode(tokenizer, "</HISTORY>\n")
+    state_close = _encode(tokenizer, "</STATE>\n")
+    question_ids = _encode(tokenizer, _section("QUESTION", question))
+    type_ids = _encode(tokenizer, _section("QUESTION_TYPE", question_type))
+    history_ids = []
+    for event in history:
+        if not isinstance(event, Mapping):
+            raise TypeError("history events must be mappings")
+        role = str(event.get("role") or "unknown")
+        label = "OBSERVATION" if role.lower() in {"tool", "observation", "environment"} else "ACTION"
+        content = _text(event.get("content"))
+        history_ids.append(_encode(tokenizer, _section(label, f"role: {role}\n{content}")))
+    option_ids = _encode(tokenizer, _option_text(option))
+    fixed = sum(map(len, (
+        state_open, system, goal, summary_open, summary_close, environment_open,
+        environment_close, open_history, close_history, state_close, question_ids,
+        type_ids, option_ids,
+    )))
+    if fixed > max_tokens:
+        raise ValueError("protected state, question, and option exceed 8192 tokens")
+    budget = max_tokens - fixed
+    summary_reserve = min(len(summary_content), max(0, budget // 3))
+    budget -= summary_reserve
+    environment_reserve = min(len(environment_content), max(0, budget // 3))
+    budget -= environment_reserve
+    selected = []
+    for event, item in reversed(list(zip(history, history_ids, strict=True))):
+        if len(item) <= budget:
+            selected.append(item)
+            budget -= len(item)
+        elif not selected and budget > 0:
+            role = str(event.get("role") or "unknown")
+            label = "OBSERVATION" if role.lower() in {"tool", "observation", "environment"} else "ACTION"
+            start = _encode(tokenizer, f"<{label}>\nrole: {role}\n")
+            end = _encode(tokenizer, f"\n</{label}>\n")
+            available = budget - len(start) - len(end)
+            if available > 0:
+                payload = _encode(tokenizer, _text(event.get("content")))
+                selected.append(start + payload[-available:] + end)
+                budget -= len(selected[-1])
+    selected.reverse()
+    extra_summary = min(budget, len(summary_content) - summary_reserve)
+    kept_summary = summary_content[:summary_reserve + extra_summary]
+    budget -= extra_summary
+    kept_environment = environment_content[:environment_reserve + budget]
+    prefix = (
+        state_open + system + goal + summary_open + kept_summary + summary_close
+        + environment_open + kept_environment + environment_close + open_history
+        + [token for item in selected for token in item] + close_history + state_close
+        + question_ids + type_ids
+    )
+    ids = prefix + option_ids
+    if len(ids) > max_tokens:
+        raise AssertionError("serialization exceeded max_tokens")
+    return ids, [False] * len(prefix) + [True] * len(option_ids)
+
+
+class _SetBlock(nn.Module):
+    def __init__(self, width: int, heads: int):
+        super().__init__()
+        self.norm1 = nn.RMSNorm(width)
+        self.attn = nn.MultiheadAttention(width, heads, dropout=0.0, batch_first=True)
+        self.norm2 = nn.RMSNorm(width)
+        self.ff = nn.Sequential(nn.Linear(width, 4 * width), nn.SiLU(),
+                                nn.Dropout(0.0), nn.Linear(4 * width, width))
+
+    def forward(self, x: torch.Tensor, valid: torch.Tensor) -> torch.Tensor:
+        normed = self.norm1(x)
+        attended, _ = self.attn(normed, normed, normed,
+                                key_padding_mask=~valid, need_weights=False)
+        x = x + attended
+        x = x + self.ff(self.norm2(x))
+        return x.masked_fill(~valid.unsqueeze(-1), 0)
+
+
+class _SetHead(nn.Module):
+    def __init__(self, width: int, layers: int, heads: int):
+        super().__init__()
+        self.input = nn.Linear(width, width)
+        self.blocks = nn.ModuleList([_SetBlock(width, heads) for _ in range(layers)])
+        self.output = nn.Sequential(nn.RMSNorm(width), nn.Linear(width, 1))
+
+    def forward(self, vectors: torch.Tensor, valid: torch.Tensor) -> torch.Tensor:
+        x = self.input(vectors.float()).masked_fill(~valid.unsqueeze(-1), 0)
+        for block in self.blocks:
+            x = block(x, valid)
+        return self.output(x).squeeze(-1).masked_fill(~valid, -torch.inf)
+
+
+class MiniJev:
+    """Score state + question + finite options, returning a probability per option."""
+
+    def __init__(self, release_dir: str | Path, device: str = "cuda"):
+        if not device.startswith("cuda") or not torch.cuda.is_available():
+            raise RuntimeError("this 4-bit BF16 release requires a CUDA GPU")
+        root = Path(release_dir)
+        config = json.loads((root / "config.json").read_text(encoding="utf-8"))
+        quant = config["quantization"]
+        if (quant["format"] != "nf4" or not quant["double_quantization"]
+                or quant["compute_dtype"] != "bfloat16" or config["pooler"] != "mean"):
+            raise ValueError("unsupported release configuration")
+        model_id = config["base_model"]
+        revision = config["base_revision"]
+        self.tokenizer = AutoTokenizer.from_pretrained(model_id, revision=revision)
+        base = AutoModel.from_pretrained(
+            model_id, revision=revision, dtype=torch.bfloat16,
+            quantization_config=BitsAndBytesConfig(
+                load_in_4bit=True, bnb_4bit_quant_type="nf4",
+                bnb_4bit_use_double_quant=True, bnb_4bit_compute_dtype=torch.bfloat16,
+            ),
+            device_map={"": device}, attn_implementation="sdpa",
+        )
+        if type(base).__name__ != "Qwen3Model":
+            raise TypeError("expected a Qwen3Model hidden-state backbone")
+        base = prepare_model_for_kbit_training(base, use_gradient_checkpointing=False)
+        self.encoder = PeftModel.from_pretrained(
+            base, root / config["adapter_dir"], is_trainable=False,
+        ).eval()
+        self.backbone = self.encoder.base_model.model
+        width = int(config["projection_dim"])
+        self.projection = nn.Linear(self.backbone.config.hidden_size, width).to(device)
+        self.scorer = _SetHead(width, int(config["set_layers"]), int(config["set_heads"])).to(device)
+        weights = load_file(root / config["head_file"], device="cpu")
+        self.projection.load_state_dict({
+            key.removeprefix("projection."): value for key, value in weights.items()
+            if key.startswith("projection.")
+        }, strict=True)
+        self.scorer.load_state_dict({
+            key.removeprefix("scorer."): value for key, value in weights.items()
+            if key.startswith("scorer.")
+        }, strict=True)
+        self.projection.eval()
+        self.scorer.eval()
+        self.device = device
+        self.max_tokens = int(config["max_tokens"])
+
+    @classmethod
+    def load(cls, release_dir: str | Path, device: str = "cuda") -> MiniJev:
+        return cls(release_dir, device)
+
+    def predict(
+        self,
+        state: Mapping[str, Any] | str,
+        question: str,
+        question_type: str,
+        answer_options: Sequence[Mapping[str, Any]],
+    ) -> dict[str, Any]:
+        if len(answer_options) < 2:
+            raise ValueError("provide at least two answer options")
+        labels = [str(option.get("label") or option.get("name") or "") for option in answer_options]
+        if any(not label for label in labels):
+            raise ValueError("every option needs a label")
+        option_ids = [str(option.get("id", i)) for i, option in enumerate(answer_options)]
+        if len(set(option_ids)) != len(option_ids):
+            raise ValueError("option IDs must be unique")
+        vectors = []
+        lengths = []
+        with torch.inference_mode():
+            for option in answer_options:
+                ids, suffix_mask = encode_branch(
+                    self.tokenizer, state, question, question_type, option, self.max_tokens,
+                )
+                input_ids = torch.tensor([ids], dtype=torch.long, device=self.device)
+                mask = torch.tensor([suffix_mask], dtype=torch.bool, device=self.device)
+                with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
+                    hidden = self.backbone(input_ids=input_ids, use_cache=False).last_hidden_state
+                pooled = (hidden.float() * mask.unsqueeze(-1)).sum(1) / mask.sum(1, keepdim=True)
+                vectors.append(self.projection(pooled).squeeze(0))
+                lengths.append(len(ids))
+            valid = torch.ones((1, len(vectors)), dtype=torch.bool, device=self.device)
+            logits = self.scorer(torch.stack(vectors).unsqueeze(0), valid)
+            probs = torch.softmax(logits.float(), dim=1)[0].cpu().tolist()
+        winner = max(range(len(probs)), key=lambda i: probs[i])
+        return {
+            "selected_id": option_ids[winner],
+            "selected_label": labels[winner],
+            "options": [
+                {"id": option_ids[i], "label": labels[i], "probability": probs[i]}
+                for i in range(len(probs))
+            ],
+            "branch_tokens": lengths,
+        }

requirements.txt

@@ -0,0 +1,6 @@
+torch==2.14.0
+transformers==5.17.0
+peft==0.21.0
+bitsandbytes==0.50.2
+safetensors==0.8.0
+accelerate==1.15.0