"""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 # installed with transformers from peft import PeftModel from safetensors.torch import load_file from transformers import AutoModelForCausalLM, AutoTokenizer, PreTrainedModel from .configuration_arbit import ArbitConfig # predict() reads at least this many tokens, whatever max_seq_length the model was loaded with. SERVE_MIN_LENGTH = 16384 # max_seq_length when none is given, as Unsloth's Clef loader defaults it. DEFAULT_LOAD_LENGTH = 4096 # Order of the per-type "temperature" list in unsloth_decision_config.json. 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.""" # --------------------------------------------------------------------------------------------------------------------- # files # --------------------------------------------------------------------------------------------------------------------- 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 # dataclasses look their module up while the file executes 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) # --------------------------------------------------------------------------------------------------------------------- # requests and answers (FastDecisionModel.predict semantics) # --------------------------------------------------------------------------------------------------------------------- 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 # --------------------------------------------------------------------------------------------------------------------- # model # --------------------------------------------------------------------------------------------------------------------- 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 # -- loading ------------------------------------------------------------------------------------------------ @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) # A requested attention implementation is the backbone's; the wrapper itself has 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, # 16-bit backbones keep the LoRA weights in 16 bits too, as Unsloth's LoRA kernels compute them. 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) # The head always runs in float32, over a backbone of any dtype. 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") # -- serving settings ------------------------------------------------------------------------------------------- @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)]) # -- forward ---------------------------------------------------------------------------------------------------- def _embedding_rows(self, weight: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: # The head averages output-embedding rows per option in its own dtype; a 16-bit backbone gets one cast copy. 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) # -- decisions -------------------------------------------------------------------------------------------------- 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 # One spare token tells a state that was cut from one that fits exactly. 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