| """Jev-style decision model following "Releasing the Inner Jev" (Geuter et al., 2026). |
| |
| A decision is rendered as a lettered multiple-choice prompt (SemIf system message, JSON user |
| message, the model's own chat template with thinking disabled). One forward pass gives the final |
| hidden state ``h`` at the last prompt position, and the decision head maps it to one logit per |
| option slot: |
| |
| z = W0 h + A std(h) + b |
| |
| ``W0`` holds the frozen language-model-head rows for the letters A..P, so the untrained model is the |
| zero-shot letter readout. ``A`` and ``b`` start at zero and are trained jointly with LoRA adapters. |
| A softmax over the question's K options (divided by a per-type temperature) gives the probabilities. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import json |
| import random |
| from collections import Counter |
| from pathlib import Path |
| from typing import Any |
|
|
| import torch |
| from torch import nn |
|
|
| SYSTEM_PROMPT = ( |
| "Apply the supplied criterion to the supplied evidence. Choose exactly one listed option. " |
| "Respond with only its uppercase letter, with no explanation or reasoning." |
| ) |
| LETTERS = "ABCDEFGHIJKLMNOP" |
| MAX_OPTIONS = len(LETTERS) |
| QUESTION_TYPES = ("choice", "noul", "score") |
| TRUNCATION_MARK = " [truncated]" |
| MASKED_LOGIT = -1e9 |
|
|
|
|
| def question_type(row: dict[str, Any]) -> str: |
| return row["question"]["type"] |
|
|
|
|
| def request_order(row: dict[str, Any]) -> list[str]: |
| """Evaluation order: choice in request order, yes/no as true then false, score low to high.""" |
| if question_type(row) == "noul": |
| return ["true", "false"] |
| return list(row["labels"]) |
|
|
|
|
| def training_order(row: dict[str, Any], rng: random.Random) -> list[str]: |
| """Options are shuffled every presentation; score levels keep their order.""" |
| order = request_order(row) |
| if question_type(row) != "score": |
| rng.shuffle(order) |
| return order |
|
|
|
|
| def user_message(row: dict[str, Any], order: list[str], state: str | None = None) -> str: |
| criteria = row["question"]["criteria"] |
| criterion = row["question"]["instructions"] |
| surface = (row.get("focus") or {}).get("surface") |
| if surface: |
| criterion = f"{criterion} Focus entity: {surface}." |
| return json.dumps({ |
| "evidence": row["state"] if state is None else state, |
| "criterion": criterion, |
| "options": [ |
| {"letter": LETTERS[index], "description": f"{label}: {criteria[label]}"} |
| for index, label in enumerate(order) |
| ], |
| }, ensure_ascii=False) |
|
|
|
|
| def _chat_ids(tokenizer: Any, user: str) -> list[int]: |
| messages = [{"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": user}] |
| text = tokenizer.apply_chat_template( |
| messages, tokenize=False, add_generation_prompt=True, enable_thinking=False |
| ) |
| return tokenizer(text, add_special_tokens=False)["input_ids"] |
|
|
|
|
| def encode( |
| tokenizer: Any, |
| row: dict[str, Any], |
| order: list[str], |
| max_length: int, |
| truncate_state: bool, |
| ) -> tuple[list[int] | None, bool]: |
| """Token ids for one decision and whether the state was truncated. |
| |
| Over-long inputs are shortened from the end of the state, never the options or the final |
| answer position. With ``truncate_state=False`` (training, as in the paper) they return None. |
| """ |
| if len(order) > MAX_OPTIONS: |
| raise ValueError(f"{row['id']}: {len(order)} options; the letter readout supports {MAX_OPTIONS}") |
| ids = _chat_ids(tokenizer, user_message(row, order)) |
| if len(ids) <= max_length: |
| return ids, False |
| if not truncate_state: |
| return None, False |
| state_ids = tokenizer(row["state"], add_special_tokens=False)["input_ids"] |
| overflow = len(ids) - max_length |
| for _ in range(4): |
| keep = len(state_ids) - overflow - 8 |
| if keep <= 0: |
| break |
| state = tokenizer.decode(state_ids[:keep]) + TRUNCATION_MARK |
| ids = _chat_ids(tokenizer, user_message(row, order, state)) |
| if len(ids) <= max_length: |
| return ids, True |
| overflow += len(ids) - max_length |
| raise ValueError(f"{row['id']}: question and options alone exceed max_length={max_length}") |
|
|
|
|
| def letter_token_ids(tokenizer: Any) -> list[int]: |
| ids = [] |
| for letter in LETTERS: |
| token_ids = tokenizer.encode(letter, add_special_tokens=False) |
| if len(token_ids) != 1: |
| raise ValueError(f"letter {letter!r} is not a single token for this tokenizer: {token_ids}") |
| ids.append(token_ids[0]) |
| return ids |
|
|
|
|
| class DecisionHead(nn.Module): |
| """z = W0 h + A std(h) + b with W0 frozen from the LM head and A, b initialised to zero.""" |
|
|
| def __init__(self, letter_rows: torch.Tensor, letter_bias: torch.Tensor | None = None, |
| softcap: float | None = None): |
| super().__init__() |
| slots, hidden = letter_rows.shape |
| self.register_buffer("letter_rows", letter_rows.detach().float().clone()) |
| self.register_buffer( |
| "letter_bias", |
| torch.zeros(slots) if letter_bias is None else letter_bias.detach().float().clone(), |
| ) |
| self.register_buffer("feature_mean", torch.zeros(hidden)) |
| self.register_buffer("feature_std", torch.ones(hidden)) |
| self.correction = nn.Parameter(torch.zeros(slots, hidden)) |
| self.bias = nn.Parameter(torch.zeros(slots)) |
| self.softcap = softcap |
|
|
| def set_feature_stats(self, features: torch.Tensor) -> None: |
| """Fix the standardisation once on training features, before training.""" |
| features = features.float() |
| self.feature_mean.copy_(features.mean(dim=0)) |
| self.feature_std.copy_(features.std(dim=0, unbiased=False).clamp(min=1e-6)) |
|
|
| def forward(self, hidden: torch.Tensor) -> torch.Tensor: |
| hidden = hidden.float() |
| readout = hidden @ self.letter_rows.T + self.letter_bias |
| if self.softcap: |
| readout = self.softcap * torch.tanh(readout / self.softcap) |
| standardized = (hidden - self.feature_mean) / self.feature_std |
| return readout + standardized @ self.correction.T + self.bias |
|
|
|
|
| class DecisionModel(nn.Module): |
| """Decoder backbone (optionally LoRA-wrapped) plus the letter-initialised decision head.""" |
|
|
| def __init__(self, backbone: nn.Module, head: DecisionHead): |
| super().__init__() |
| self.backbone = backbone |
| self.head = head |
|
|
| def last_hidden(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor: |
| hidden = self.backbone( |
| input_ids=input_ids, attention_mask=attention_mask, use_cache=False |
| ).last_hidden_state |
| last = attention_mask.sum(dim=1).clamp(min=1) - 1 |
| pooled = hidden[torch.arange(hidden.size(0), device=hidden.device), last.to(hidden.device)] |
| return pooled.to(self.head.letter_rows.device) |
|
|
| def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor, |
| option_counts: torch.Tensor) -> torch.Tensor: |
| logits = self.head(self.last_hidden(input_ids, attention_mask)) |
| slots = torch.arange(logits.size(1), device=logits.device) |
| invalid = slots.unsqueeze(0) >= option_counts.to(logits.device).unsqueeze(1) |
| return logits.masked_fill(invalid, MASKED_LOGIT) |
|
|
|
|
| def pad_batch(sequences: list[list[int]], pad_id: int, device: torch.device | str): |
| width = max(len(sequence) for sequence in sequences) |
| input_ids = torch.full((len(sequences), width), pad_id, dtype=torch.long) |
| attention_mask = torch.zeros((len(sequences), width), dtype=torch.long) |
| for index, sequence in enumerate(sequences): |
| input_ids[index, :len(sequence)] = torch.tensor(sequence, dtype=torch.long) |
| attention_mask[index, :len(sequence)] = 1 |
| return input_ids.to(device), attention_mask.to(device) |
|
|
|
|
| def parse_max_memory(spec: str | None) -> dict[Any, str] | None: |
| if not spec: |
| return None |
| result: dict[Any, str] = {} |
| for part in spec.split(","): |
| key, value = part.split("=") |
| result[int(key) if key.strip().isdigit() else key.strip()] = value.strip() |
| return result |
|
|
|
|
| def pick_dtype(counter: Counter) -> torch.dtype: |
| if not torch.cuda.is_available(): |
| counter["cpu_execution_float32"] += 1 |
| return torch.float32 |
| if all(torch.cuda.get_device_capability(index)[0] >= 8 for index in range(torch.cuda.device_count())): |
| return torch.bfloat16 |
| counter["float16_no_bf16_device"] += 1 |
| return torch.float16 |
|
|
|
|
| def load_decision_model( |
| model_id: str, |
| *, |
| quantization: str = "none", |
| max_memory: str | None = None, |
| adapter_dir: Path | None = None, |
| head_path: Path | None = None, |
| trainable_adapter: bool = False, |
| counter: Counter | None = None, |
| local_files_only: bool = False, |
| ): |
| """Load backbone, tokenizer and head. Returns (model, tokenizer, letter_ids, dtype).""" |
| from transformers import AutoModelForCausalLM, AutoTokenizer |
|
|
| counter = counter if counter is not None else Counter() |
| tokenizer = AutoTokenizer.from_pretrained(model_id, local_files_only=local_files_only) |
| if tokenizer.pad_token_id is None: |
| tokenizer.pad_token = tokenizer.eos_token |
| counter["pad_token_from_eos"] += 1 |
| dtype = pick_dtype(counter) |
| if torch.cuda.is_available() and any( |
| torch.cuda.get_device_capability(index)[0] < 8 for index in range(torch.cuda.device_count()) |
| ): |
| |
| torch.backends.cuda.enable_flash_sdp(False) |
| counter["flash_sdp_disabled_pre_ampere_gpu"] += 1 |
| kwargs: dict[str, Any] = {"dtype": dtype, "local_files_only": local_files_only, |
| "attn_implementation": "sdpa"} |
| if torch.cuda.is_available(): |
| kwargs["device_map"] = "auto" |
| memory = parse_max_memory(max_memory) |
| if memory: |
| kwargs["max_memory"] = memory |
| if quantization == "4bit": |
| from transformers import BitsAndBytesConfig |
| kwargs["quantization_config"] = BitsAndBytesConfig( |
| load_in_4bit=True, |
| bnb_4bit_quant_type="nf4", |
| bnb_4bit_use_double_quant=True, |
| bnb_4bit_compute_dtype=dtype, |
| ) |
| elif quantization != "none": |
| raise ValueError(f"unknown quantization {quantization!r}") |
|
|
| language_model = AutoModelForCausalLM.from_pretrained(model_id, **kwargs) |
| letter_ids = letter_token_ids(tokenizer) |
| offloaded = sorted({str(d) for d in (getattr(language_model, "hf_device_map", None) or {}).values() |
| if str(d) in ("cpu", "disk")}) |
| if offloaded: |
| counter["layers_offloaded_to_" + "_".join(offloaded)] += 1 |
| output = language_model.get_output_embeddings() |
| weight = output.weight |
| if weight.device.type == "meta": |
| |
| weight = output._hf_hook.weights_map["weight"] |
| letter_rows = weight[letter_ids].detach().float().cpu() |
| letter_bias = output.bias[letter_ids].detach().float().cpu() if getattr(output, "bias", None) is not None else None |
| config = language_model.config |
| softcap = getattr(config, "final_logit_softcapping", None) or getattr( |
| getattr(config, "text_config", None), "final_logit_softcapping", None |
| ) |
| backbone = language_model.get_decoder() |
| del language_model |
|
|
| head = DecisionHead(letter_rows, letter_bias, softcap) |
| if head_path is not None: |
| payload = torch.load(head_path, map_location="cpu", weights_only=True) |
| if payload["letter_ids"] != letter_ids: |
| raise ValueError("checkpoint letter ids do not match this tokenizer") |
| head.load_state_dict(payload["head_state_dict"]) |
| head_device = next(backbone.parameters()).device |
| head.to(head_device) |
|
|
| if adapter_dir is not None: |
| from peft import PeftModel |
| backbone = PeftModel.from_pretrained(backbone, str(adapter_dir), is_trainable=trainable_adapter) |
| model = DecisionModel(backbone, head) |
| return model, tokenizer, letter_ids, dtype |
|
|
|
|
| def option_logits(model: DecisionModel, tokenizer: Any, row: dict[str, Any], max_length: int, |
| counter: Counter) -> tuple[dict[str, float], bool]: |
| """Raw decision-head logits for the row's options, in request order (one forward pass).""" |
| order = request_order(row) |
| ids, truncated = encode(tokenizer, row, order, max_length, truncate_state=True) |
| if truncated: |
| counter["state_truncated"] += 1 |
| device = next(model.backbone.parameters()).device |
| input_ids, attention_mask = pad_batch([ids], tokenizer.pad_token_id, device) |
| counts = torch.tensor([len(order)]) |
| model.eval() |
| with torch.inference_mode(): |
| logits = model(input_ids, attention_mask, counts)[0, :len(order)].float().cpu() |
| return {label: float(logits[index]) for index, label in enumerate(order)}, truncated |
|
|
|
|
| def probabilities(logits: dict[str, float], temperature: float = 1.0) -> dict[str, float]: |
| values = torch.tensor(list(logits.values()), dtype=torch.float64) / temperature |
| probs = torch.softmax(values, dim=0).tolist() |
| return dict(zip(logits, probs)) |
|
|
|
|
| def jev_answer(row: dict[str, Any], probs: dict[str, float]) -> dict[str, Any]: |
| """Typed answer in TypeSafe's format: choice, noul (P(true)) or score.""" |
| kind = question_type(row) |
| best = max(probs, key=probs.get) |
| if kind == "noul": |
| return {"type": "noul", "noul": probs["true"]} |
| if kind == "score": |
| return {"type": "score", "score": int(best) if best.isdigit() else best, "probabilities": probs} |
| return {"type": "choice", "choice": best, "probabilities": probs} |
|
|
