"""Tiny-Jev — a small "System One" decision model: typed Choice / Score / Noul answers with calibrated probabilities. from transformers import AutoModel, AutoTokenizer tok = AutoTokenizer.from_pretrained("lostargon/Tiny-Jev") model = AutoModel.from_pretrained("lostargon/Tiny-Jev", trust_remote_code=True).eval() model.choice(tok, "My card was charged twice and nobody answers.", "Which team should handle this", ["billing", "technical", "sales"]) # {'choice': 'billing', 'probabilities': {'billing': 0.97, 'technical': 0.02, 'sales': 0.01}, 'confidence': 0.97} model.noul(tok, "Customer: I want my money back now.", "The customer requests a refund") # 0.98 model.score(tok, "Arrived on time, works great.", "How satisfied is the reviewer", ["very unhappy", "neutral", "very happy"]) model.decide(tok, state, [{"kind": "choice", "instructions": ..., "criteria": {...}}, {"kind": "noul", "instructions": ...}]) The whole model — decoder stack and decision head — lives in one safetensors file. No LoRA, no generation. """ from __future__ import annotations import json import torch import torch.nn as nn import torch.nn.functional as F from transformers import AutoConfig, AutoModel from transformers.models.qwen3.configuration_qwen3 import Qwen3Config from transformers.models.qwen3.modeling_qwen3 import Qwen3Model, Qwen3PreTrainedModel class TinyJevConfig(Qwen3Config): model_type = "tiny_jev" def __init__(self, temperature: float = 1.0, opt_token: str = "", cue: str = " - correct?", **kwargs): super().__init__(**kwargs) self.temperature = temperature self.opt_token = opt_token self.cue = cue class TinyJevModel(Qwen3PreTrainedModel): config_class = TinyJevConfig def __init__(self, config: TinyJevConfig): super().__init__(config) self.model = Qwen3Model(config) self.head = nn.Linear(config.hidden_size, 1) self.post_init() # ---- core --------------------------------------------------------------------------------------------- def forward(self, input_ids, attention_mask, opt_positions, opt_mask=None): """opt_positions: [B, O] token indices of the option markers; returns logits [B, O] (masked with -inf).""" h = self.model(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state g = torch.gather(h, 1, opt_positions.unsqueeze(-1).expand(-1, -1, h.shape[-1])) g = F.layer_norm(g.float(), (h.shape[-1],)).to(self.head.weight.dtype) logits = self.head(g).float().squeeze(-1) / self.config.temperature if opt_mask is not None: logits = logits.masked_fill(~opt_mask, float("-inf")) return logits # ---- prompt building --------------------------------------------------------------------------------- def encode(self, tok, state: str, question: str, options: list[str], max_len: int = 4096): ids_of = lambda t: tok(t, add_special_tokens=False)["input_ids"] opt_id = tok.convert_tokens_to_ids(self.config.opt_token) pre = ids_of("\n") mid = ids_of("\n\n\n" + question + "\n\n\n") opts = [ids_of(o + self.config.cue) + [opt_id] + ids_of("\n") for o in options] post = ids_of("") budget = max_len - (len(pre) + len(mid) + sum(map(len, opts)) + len(post)) st = ids_of(state) if len(st) > budget: # truncate the state only, keeping head and tail head = max(0, int(budget * 0.7)) st = st[:head] + st[-(budget - head):] if budget > 0 else [] ids, pos = pre + st + mid, [] for o in opts: pos.append(len(ids) + len(o) - 3) # the cue's last token, right before ids += o return ids + post, pos @torch.no_grad() def probabilities(self, tok, state, questions: list[tuple[str, list[str]]]) -> list[list[float]]: state = state if isinstance(state, str) else json.dumps(state, ensure_ascii=False) encs = [self.encode(tok, state, q, o) for q, o in questions] L, O = max(len(e[0]) for e in encs), max(len(e[1]) for e in encs) pad = tok.pad_token_id if tok.pad_token_id is not None else 0 dev = next(self.parameters()).device ids = torch.full((len(encs), L), pad, dtype=torch.long); att = torch.zeros_like(ids) pos = torch.zeros((len(encs), O), dtype=torch.long); mask = torch.zeros((len(encs), O), dtype=torch.bool) for i, (e, p) in enumerate(encs): ids[i, :len(e)] = torch.tensor(e); att[i, :len(e)] = 1 pos[i, :len(p)] = torch.tensor(p); mask[i, :len(p)] = True logits = self(ids.to(dev), att.to(dev), pos.to(dev), mask.to(dev)) probs = F.softmax(logits, -1).cpu() return [probs[i, :len(p)].tolist() for i, (_, p) in enumerate(encs)] # ---- typed API ----------------------------------------------------------------------------------------- def decide(self, tok, state, questions: list[dict]) -> list[dict]: specs, keys = [], [] for q in questions: crit = q.get("criteria") if q["kind"] == "noul": specs.append((q["instructions"], ["no", "yes"])); keys.append(None) elif isinstance(crit, dict): specs.append((q["instructions"], [f"{k}: {v}" for k, v in crit.items()])); keys.append(list(crit)) else: specs.append((q["instructions"], list(crit))); keys.append(list(crit)) out = [] for q, k, p in zip(questions, keys, self.probabilities(tok, state, specs)): if q["kind"] == "noul": out.append({"noul": p[1]}) elif q["kind"] == "choice": i = max(range(len(p)), key=p.__getitem__) out.append({"choice": k[i], "probabilities": dict(zip(k, p)), "confidence": p[i]}) else: out.append({"score": sum(i * x for i, x in enumerate(p)), "probabilities": dict(zip(k, p)), "confidence": max(p)}) return out def choice(self, tok, state, instructions, criteria): return self.decide(tok, state, [{"kind": "choice", "instructions": instructions, "criteria": criteria}])[0] def score(self, tok, state, instructions, levels): return self.decide(tok, state, [{"kind": "score", "instructions": instructions, "criteria": levels}])[0] def noul(self, tok, state, instructions): return self.decide(tok, state, [{"kind": "noul", "instructions": instructions}])[0]["noul"] AutoConfig.register("tiny_jev", TinyJevConfig) AutoModel.register(TinyJevConfig, TinyJevModel)