| """Arbit: a typed-decision model for stock Transformers + PEFT. |
| |
| from transformers import AutoModel |
| |
| model = AutoModel.from_pretrained("litillabs/arbit-1.7b", trust_remote_code=True) |
| answers = model.predict(state, questions) # {name: answer} |
| body = model.systemone({"model": "arbit-1.7b", "state": state, "questions": questions}) |
| |
| Arbit is Qwen3-1.7B with a LoRA adapter and a joint schema head of Cloudflare's Clef design. Loading downloads the |
| pinned base (``config.base_model`` at ``config.base_revision``), applies the adapter and loads the head from |
| ``config.checkpoint_subfolder``, and imports Cloudflare's reference code (``clef/joint_schema_model.py``, vendored |
| unmodified) for the record encoding, the head modules and the Decision API answers. |
| |
| Serving follows Unsloth's ``FastDecisionModel.predict`` for this checkpoint layout: |
| |
| * a JSON object or array given as a string state is decoded first; questions are normalised (a ``choice`` list of |
| labels becomes options without descriptions) and checked; |
| * the text-only record is encoded by Clef's ``encode_record`` and read up to ``max(max_seq_length, 16384)`` tokens |
| (past that the end of the state is cut, never the questions or options); |
| * the backbone's last hidden state feeds the head, which runs in float32 whatever the backbone dtype; |
| * each question's logits are divided by its served temperature (``head_temperature``, 1 once folded into the head |
| weights, times the per-type ``temperature`` in ``unsloth_decision_config.json``, ordered choice, score, noul) |
| before the softmax; |
| * each answer is Clef's Decision API answer plus ``"answer"`` (the option, True / False for noul, the level number |
| for score) and ``"probabilities"`` over every option, rounded to 4 decimals. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import hashlib |
| import importlib.util |
| import json |
| import math |
| import sys |
| import warnings |
| from pathlib import Path |
| from types import ModuleType |
| from typing import Any |
|
|
| import torch |
| import transformers |
| from huggingface_hub import hf_hub_download |
| from peft import PeftModel |
| from safetensors.torch import load_file |
| from transformers import AutoModelForCausalLM, AutoTokenizer, PreTrainedModel |
|
|
| from .configuration_arbit import ArbitConfig |
|
|
| |
| SERVE_MIN_LENGTH = 16384 |
| |
| DEFAULT_LOAD_LENGTH = 4096 |
| |
| TEMPERATURE_TYPES = ("choice", "score", "noul") |
| QUESTION_KINDS = frozenset(TEMPERATURE_TYPES) |
|
|
| _UNSET = object() |
| _LOADER_ONLY_KWARGS = ( |
| "_from_auto", |
| "_fast_init", |
| "adapter_kwargs", |
| "code_revision", |
| "ignore_mismatched_sizes", |
| "output_loading_info", |
| "proxies", |
| "resume_download", |
| "state_dict", |
| "subfolder", |
| "trust_remote_code", |
| "use_auth_token", |
| "use_safetensors", |
| "variant", |
| "weights_only", |
| ) |
|
|
|
|
| class DecisionInputError(ValueError): |
| """A request the model cannot answer as given.""" |
|
|
|
|
| |
| |
| |
| def _resolve_file(source: str, filename: str, *, revision: str | None, hub: dict[str, Any]) -> Path: |
| """A file of the model repository: from the folder when source is a directory, else from the Hub cache.""" |
| root = Path(source) |
| if root.is_dir(): |
| path = root / filename |
| if not path.is_file(): |
| raise FileNotFoundError(f"{path} is missing from the Arbit folder") |
| return path |
| return Path(hf_hub_download(repo_id=source, filename=filename, revision=revision, **hub)) |
|
|
|
|
| def _load_clef(path: Path, expected_sha256: str | None) -> ModuleType: |
| """Import the vendored Clef reference module by path (Transformers' remote-code loader copies only flat files).""" |
| data = path.read_bytes() |
| digest = hashlib.sha256(data).hexdigest() |
| if expected_sha256 and digest != expected_sha256: |
| raise ValueError( |
| f"{path} has sha256 {digest}, not the vendored Clef reference {expected_sha256}; " |
| "restore the unmodified file (or set clef_module_sha256 to null in config.json)" |
| ) |
| name = f"arbit_clef_joint_schema_model_{digest[:16]}" |
| module = sys.modules.get(name) |
| if module is None: |
| spec = importlib.util.spec_from_file_location(name, path) |
| if spec is None or spec.loader is None: |
| raise ImportError(f"cannot import {path}") |
| module = importlib.util.module_from_spec(spec) |
| sys.modules[name] = module |
| try: |
| spec.loader.exec_module(module) |
| except BaseException: |
| sys.modules.pop(name, None) |
| raise |
| return module |
|
|
|
|
| def _check_decision_config(decision_config: dict[str, Any], base_model: str) -> None: |
| if decision_config.get("layout") != "clef": |
| raise ValueError(f"unsupported decision layout {decision_config.get('layout')!r}; expected 'clef'") |
| temperatures = decision_config.get("temperature", [1.0] * 3) |
| if len(temperatures) != 3 or not all( |
| isinstance(t, (int, float)) and math.isfinite(t) and t > 0 for t in temperatures |
| ): |
| raise ValueError(f"temperature must be three positive numbers (choice, score, noul), got {temperatures!r}") |
| head_temperature = decision_config.get("head_temperature", 1.0) |
| if not (isinstance(head_temperature, (int, float)) and math.isfinite(head_temperature) and head_temperature > 0): |
| raise ValueError(f"head_temperature must be a positive number, got {head_temperature!r}") |
| if decision_config.get("temperature_by_options"): |
| raise NotImplementedError("per-option-count temperatures (temperature_by_options) are not supported") |
| if decision_config.get("load_in_4bit"): |
| warnings.warn("this checkpoint was trained over a 4-bit base; it is loaded here in 16 or 32 bits") |
| if decision_config.get("base_model") not in (None, base_model): |
| warnings.warn( |
| f"config.json names base {base_model!r} but the checkpoint was trained from " |
| f"{decision_config['base_model']!r}" |
| ) |
|
|
|
|
| def _torch_dtype(value: Any, device: torch.device) -> torch.dtype: |
| if value is None: |
| return torch.float32 |
| if value == "auto": |
| if device.type == "cuda" and torch.cuda.is_bf16_supported(): |
| return torch.bfloat16 |
| return torch.float32 |
| if isinstance(value, str): |
| value = getattr(torch, value.replace("torch.", ""), None) |
| if value not in (torch.float32, torch.bfloat16, torch.float16): |
| raise ValueError(f"dtype must be float32, bfloat16, float16 or 'auto', got {value!r}") |
| return value |
|
|
|
|
| def _pick_device(device: Any, device_map: Any) -> torch.device: |
| if device is None and device_map is not None: |
| if isinstance(device_map, dict): |
| if set(device_map) != {""}: |
| raise ValueError('Arbit loads onto one device: pass device="cuda:0" (or device_map={"": ...})') |
| device = device_map[""] |
| else: |
| device = device_map |
| if device in (None, "auto"): |
| if torch.cuda.is_available(): |
| return torch.device("cuda") |
| if torch.backends.mps.is_available(): |
| return torch.device("mps") |
| return torch.device("cpu") |
| if isinstance(device, int): |
| return torch.device("cuda", device) |
| return torch.device(device) |
|
|
|
|
| |
| |
| |
| def parse_state(state: Any) -> Any: |
| """A JSON object or array passed as a string is decoded; anything else is used as given.""" |
| if isinstance(state, str) and state.strip()[:1] in ("{", "["): |
| try: |
| return json.loads(state) |
| except ValueError: |
| return state |
| return state |
|
|
|
|
| def normalize_question(name: str, question: Any) -> dict[str, Any]: |
| """The question as the encoder reads it: type, instructions and (normalised) criteria only.""" |
| if not isinstance(question, dict) or question.get("type") not in QUESTION_KINDS: |
| raise DecisionInputError(f'"{name}" is not a valid question (type must be choice, score or noul)') |
| kind, criteria = question["type"], question.get("criteria") |
| if kind == "choice": |
| if isinstance(criteria, list) and criteria and all(isinstance(label, str) for label in criteria): |
| criteria = dict.fromkeys(criteria) |
| if not (isinstance(criteria, dict) and criteria): |
| raise DecisionInputError(f'"{name}" needs criteria naming its options') |
| if len({str(key) for key in criteria}) != len(criteria): |
| raise DecisionInputError(f'"{name}" has repeated options') |
| elif kind == "score": |
| if not (isinstance(criteria, list) and criteria): |
| raise DecisionInputError(f'"{name}" needs a list of criteria levels') |
| elif criteria is not None and (not isinstance(criteria, dict) or set(criteria) - {"true", "false"}): |
| raise DecisionInputError(f'"{name}" criteria may only have "true" and "false"') |
| normalized = {"type": kind, "instructions": question.get("instructions")} |
| if criteria is not None: |
| normalized["criteria"] = criteria |
| return normalized |
|
|
|
|
| def with_prediction(question: dict[str, Any], answer: dict[str, Any], probabilities: dict[str, float]) -> dict: |
| """Decision API answer plus "answer" and 4-decimal "probabilities" (kept from the answer when it has them).""" |
| kind = question["type"] |
| best = max(probabilities, key=probabilities.__getitem__) |
| out = dict(answer) |
| out["answer"] = int(best) if kind == "score" else (best == "true" if kind == "noul" else best) |
| out["probabilities"] = answer.get("probabilities") or { |
| key: round(float(value), 4) for key, value in probabilities.items() |
| } |
| return out |
|
|
|
|
| |
| |
| |
| class ArbitDecisionModel(PreTrainedModel): |
| """Qwen3 backbone with LoRA (PEFT) and a Clef joint schema head; answers typed decisions in one forward pass.""" |
|
|
| config_class = ArbitConfig |
| base_model_prefix = "language_model" |
| main_input_name = "input_ids" |
|
|
| def __init__( |
| self, |
| config: ArbitConfig, |
| *, |
| language_model: torch.nn.Module | None = None, |
| head: torch.nn.Module | None = None, |
| tokenizer: Any = None, |
| clef: ModuleType | None = None, |
| decision_config: dict[str, Any] | None = None, |
| ) -> None: |
| super().__init__(config) |
| if language_model is None or head is None or tokenizer is None or clef is None: |
| raise ValueError("load Arbit with ArbitDecisionModel.from_pretrained() or AutoModel.from_pretrained()") |
| self.language_model = language_model |
| self.head = head |
| self.tokenizer = tokenizer |
| self.clef = clef |
| self.decision_config = dict(decision_config or {}) |
| self._lexical_table: torch.Tensor | None = None |
|
|
| |
| @classmethod |
| def from_pretrained( |
| cls, |
| pretrained_model_name_or_path: str | Path, |
| *model_args: Any, |
| config: ArbitConfig | None = None, |
| revision: str | None = None, |
| token: str | bool | None = None, |
| cache_dir: str | Path | None = None, |
| force_download: bool = False, |
| local_files_only: bool = False, |
| dtype: Any = None, |
| torch_dtype: Any = None, |
| device: Any = None, |
| device_map: Any = None, |
| max_seq_length: Any = _UNSET, |
| attn_implementation: str | None = None, |
| **kwargs: Any, |
| ) -> "ArbitDecisionModel": |
| """Load the pinned base, apply the LoRA adapter, load the Clef head; float32 unless dtype says otherwise. |
| |
| dtype: None (float32), "auto" (bfloat16 on a CUDA GPU that supports it, else float32), or a torch dtype. |
| device: "cuda", "mps", "cpu", ... (default: CUDA, then MPS, then CPU). max_seq_length: as Unsloth's. |
| Other keyword arguments go to the base model's from_pretrained. |
| """ |
| if model_args: |
| raise TypeError("ArbitDecisionModel.from_pretrained takes keyword arguments only") |
| source = str(pretrained_model_name_or_path) |
| hub = { |
| "token": token, |
| "cache_dir": cache_dir, |
| "force_download": force_download, |
| "local_files_only": local_files_only, |
| } |
| if config is None: |
| config = ArbitConfig.from_pretrained(source, revision=revision, **hub) |
| if max_seq_length is not _UNSET: |
| config.max_seq_length = max_seq_length |
| artifact_revision = revision or getattr(config, "_commit_hash", None) |
| |
| attn_implementation = attn_implementation or getattr(config, "_attn_implementation_internal", None) |
| config._attn_implementation_internal = None |
| base_kwargs = {key: value for key, value in kwargs.items() if key not in _LOADER_ONLY_KWARGS} |
|
|
| target = _pick_device(device, device_map) |
| weight_dtype = _torch_dtype(dtype if dtype is not None else torch_dtype, target) |
| if target.type == "mps" and weight_dtype != torch.float32: |
| warnings.warn("16-bit backbones on MPS have given wrong answers on some macOS releases; float32 is verified") |
|
|
| subfolder = (config.checkpoint_subfolder or "").strip("/") |
|
|
| def checkpoint_file(name: str) -> Path: |
| return _resolve_file(source, f"{subfolder}/{name}" if subfolder else name, revision=artifact_revision, hub=hub) |
|
|
| clef = _load_clef( |
| _resolve_file(source, config.clef_module_file, revision=artifact_revision, hub=hub), |
| config.clef_module_sha256, |
| ) |
| decision_config = json.loads(checkpoint_file(config.decision_config_file).read_text(encoding="utf-8")) |
| _check_decision_config(decision_config, config.base_model) |
| head_config = json.loads(checkpoint_file(config.head_config_file).read_text(encoding="utf-8")) |
|
|
| local = Path(source).is_dir() |
| checkpoint_source = str(Path(source) / subfolder) if local else source |
| checkpoint_subfolder = None if local or not subfolder else subfolder |
| tokenizer = AutoTokenizer.from_pretrained( |
| checkpoint_source, subfolder=checkpoint_subfolder or "", revision=artifact_revision, **hub |
| ) |
|
|
| dtype_key = "dtype" if _transformers_at_least(4, 56) else "torch_dtype" |
| base = AutoModelForCausalLM.from_pretrained( |
| config.base_model, |
| revision=config.base_revision, |
| attn_implementation=attn_implementation or "sdpa", |
| **{dtype_key: weight_dtype}, |
| **hub, |
| **base_kwargs, |
| ) |
| base.config.use_cache = False |
| language_model = PeftModel.from_pretrained( |
| base, |
| checkpoint_source, |
| subfolder=checkpoint_subfolder, |
| revision=artifact_revision, |
| is_trainable=False, |
| |
| autocast_adapter_dtype=weight_dtype == torch.float32, |
| **hub, |
| ) |
| language_model.to(target) |
|
|
| head = clef.JointSchemaHead(**head_config) |
| head.load_state_dict(load_file(str(checkpoint_file(config.head_weights_file))), strict=True) |
| |
| head = head.to(device=target, dtype=torch.float32) |
|
|
| model = cls( |
| config, |
| language_model=language_model, |
| head=head, |
| tokenizer=tokenizer, |
| clef=clef, |
| decision_config=decision_config, |
| ) |
| model.eval() |
| return model |
|
|
| def save_pretrained(self, *args: Any, **kwargs: Any) -> None: |
| raise NotImplementedError("Arbit is published as files; copy the model folder instead of re-saving it") |
|
|
| |
| @property |
| def serve_max_length(self) -> int: |
| """Tokens predict() reads: max(max_seq_length, 16384), with max_seq_length 4096 when unset.""" |
| return max(int(self.config.max_seq_length or DEFAULT_LOAD_LENGTH), SERVE_MIN_LENGTH) |
|
|
| def served_temperature(self, question_type: str) -> float: |
| """head_temperature (1 once folded into the head) times the per-type temperature.""" |
| temperatures = self.decision_config.get("temperature", [1.0] * 3) |
| head_temperature = self.decision_config.get("head_temperature", 1.0) |
| return float(head_temperature) * float(temperatures[TEMPERATURE_TYPES.index(question_type)]) |
|
|
| |
| def _embedding_rows(self, weight: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: |
| |
| if weight.dtype == dtype: |
| return weight |
| cached = self._lexical_table |
| if cached is None or cached.device != weight.device or cached.shape != weight.shape: |
| cached = weight.detach().to(dtype) |
| self._lexical_table = cached |
| return cached |
|
|
| def forward( |
| self, |
| input_ids: torch.Tensor, |
| attention_mask: torch.Tensor, |
| records: list[Any], |
| **_: Any, |
| ) -> list[list[torch.Tensor]]: |
| """Raw joint-head logits, before temperature: per record, one tensor of option logits per question. |
| |
| Takes the batch from the Clef reference ``collate_records`` (input_ids, attention_mask, records). |
| """ |
| backbone = self.language_model.get_base_model() |
| hidden = backbone.model( |
| input_ids=input_ids, |
| attention_mask=attention_mask, |
| use_cache=False, |
| return_dict=True, |
| ).last_hidden_state |
| dtype = self.head.hidden_norm.weight.dtype |
| rows = self._embedding_rows(backbone.get_output_embeddings().weight, dtype) |
| return self.head(hidden.to(dtype), input_ids, attention_mask, records, rows) |
|
|
| |
| def _encode(self, state: Any, questions: dict[str, dict[str, Any]], max_length: int) -> Any: |
| return self.clef.encode_record(self.tokenizer, {"state": state, "questions": questions}, max_length=max_length) |
|
|
| def _decide(self, state: Any, questions: dict[str, dict[str, Any]], predicted: bool) -> dict[str, Any]: |
| max_length = self.serve_max_length |
| encoded = self._encode(state, questions, max_length) |
| batch = self.clef.collate_records([encoded], self.tokenizer.pad_token_id, self.head.hidden_norm.weight.device) |
| if self.training: |
| self.eval() |
| with torch.inference_mode(): |
| logits = self(batch["input_ids"], batch["attention_mask"], batch["records"])[0] |
| answers: dict[str, Any] = {} |
| distributions: dict[str, dict[str, float]] = {} |
| for question, row in zip(encoded.questions, logits): |
| asked = questions[question.question_id] |
| scaled = row.float().cpu() / self.served_temperature(asked["type"]) |
| probabilities = dict(zip(question.option_ids, scaled.softmax(-1).tolist())) |
| answer = self.clef.systemone_answer(asked, probabilities) |
| answers[question.question_id] = with_prediction(asked, answer, probabilities) if predicted else answer |
| distributions[question.question_id] = probabilities |
| |
| truncated = len(encoded.input_ids) >= max_length and ( |
| len(self._encode(state, questions, max_length + 1).input_ids) > max_length |
| ) |
| return { |
| "answers": answers, |
| "probabilities": distributions, |
| "input_tokens": len(encoded.input_ids), |
| "max_length": max_length, |
| "truncated": truncated, |
| } |
|
|
| def decide(self, state: Any, questions: dict[str, Any]) -> dict[str, Any]: |
| """predict() plus the unrounded per-option probabilities, the input length and whether the state was cut.""" |
| if not isinstance(questions, dict) or not questions: |
| raise DecisionInputError("questions must be a non-empty dict of name to question") |
| asked = {str(name): normalize_question(str(name), question) for name, question in questions.items()} |
| return self._decide(parse_state(state), asked, predicted=True) |
|
|
| def count_tokens(self, state: Any, questions: dict[str, Any], *, decode_state: bool = True) -> int: |
| """Tokens the whole input takes before any cut; predict() and systemone() read at most serve_max_length. |
| |
| decode_state=True counts the input as predict() reads it (a JSON string state is decoded first); |
| False counts it as systemone() reads it (the state as given).""" |
| if not isinstance(questions, dict) or not questions: |
| raise DecisionInputError("questions must be a non-empty dict of name to question") |
| asked = {str(name): normalize_question(str(name), question) for name, question in questions.items()} |
| return len(self._encode(parse_state(state) if decode_state else state, asked, 10**9).input_ids) |
|
|
| def predict(self, state: Any, questions: dict[str, Any]) -> dict[str, Any]: |
| """Answer one Decision API request: {name: answer} at the calibrated temperatures. |
| |
| Each answer is the Decision API answer (choice / confidence, score / legend, or noul) plus "answer" (the |
| option, True / False for noul, the level number for score) and "probabilities" over every option. |
| """ |
| return self.decide(state, questions)["answers"] |
|
|
| def systemone(self, request: dict[str, Any]) -> dict[str, Any]: |
| """Answer a Jev / SystemOne ``/v1/systemone`` request body with its response body. |
| |
| The request has ``model``, ``state`` (any JSON value, used as given) and ``questions``. |
| """ |
| if not isinstance(request, dict): |
| raise DecisionInputError("the request must be a JSON object") |
| questions = request.get("questions") |
| if not isinstance(request.get("model"), str) or "state" not in request: |
| raise DecisionInputError("model and state are required") |
| if not isinstance(questions, dict) or not questions: |
| raise DecisionInputError("at least one question is required") |
| for question_id, question in questions.items(): |
| if not isinstance(question, dict) or question.get("type") not in QUESTION_KINDS: |
| raise DecisionInputError(f"{question_id}: type must be noul, choice, or score") |
| if question["type"] != "noul" and not question.get("criteria"): |
| raise DecisionInputError(f"{question_id}: criteria must not be empty") |
| if request.get("images") or request.get("videos"): |
| raise DecisionInputError("Arbit reads text only; images and videos are not supported") |
| asked = {str(name): normalize_question(str(name), question) for name, question in questions.items()} |
| result = self._decide(request["state"], asked, predicted=False) |
| return { |
| "model": request["model"], |
| "answers": result["answers"], |
| "usage": {"input_tokens": result["input_tokens"], "output_tokens": 0}, |
| } |
|
|
|
|
| def _transformers_at_least(major: int, minor: int) -> bool: |
| parts = transformers.__version__.split(".") |
| try: |
| return (int(parts[0]), int(parts[1])) >= (major, minor) |
| except (IndexError, ValueError): |
| return True |
|
|
