| """Arbiter: load the model from this repository and answer typed decisions. |
| |
| from arbiter import Arbiter |
| from questions import OUTLOOK |
| |
| arbiter = Arbiter.from_pretrained("ternary/Arbiter-8B") |
| answer = arbiter.decide(state="...article or 10-K passage...", question=OUTLOOK, focus="Meta Platforms") |
| print(answer["choice"], answer["probabilities"]) |
| |
| One forward pass per decision; the answer carries a probability for every option, calibrated with the |
| per-type temperatures in calibration.json. Every fallback (state truncated to fit, uncalibrated |
| temperature) is counted in `arbiter.fallback_counts`. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import json |
| from collections import Counter |
| from pathlib import Path |
| from typing import Any |
|
|
| from model import QUESTION_TYPES, jev_answer, load_decision_model, option_logits, probabilities, question_type |
|
|
|
|
| class Arbiter: |
| def __init__(self, model, tokenizer, temperatures: dict[str, float], max_length: int, counter: Counter): |
| self.model, self.tokenizer = model, tokenizer |
| self.temperatures, self.max_length = temperatures, max_length |
| self.fallback_counts = counter |
|
|
| @classmethod |
| def from_pretrained(cls, path_or_repo: str = "ternary/Arbiter-8B", adapter: bool = False, |
| quantization: str | None = None, max_memory: str | None = None) -> "Arbiter": |
| """Load from a local folder or a Hugging Face repo id. |
| Default: the full model (LoRA merged into the weights, fp16/bf16, no base download). |
| adapter=True: the LoRA in adapter/ on Qwen/Qwen3-8B, 4-bit as trained (6.5 GB GPU memory measured). |
| quantization overrides the default ('none' or '4bit').""" |
| from safetensors.torch import load_file |
|
|
| path = Path(path_or_repo) |
| if not path.exists(): |
| from huggingface_hub import snapshot_download |
| path = Path(snapshot_download(path_or_repo)) |
| meta = json.loads((path / "metadata.json").read_text()) |
| head_config = json.loads((path / "head_config.json").read_text()) |
| counter: Counter = Counter() |
| if (path / "config.json").exists() and not adapter: |
| |
| model_id, adapter_dir, default_quantization = str(path), None, "none" |
| calibration = path / "calibration.json" |
| else: |
| |
| model_id, adapter_dir, default_quantization = meta["base_model"], path / "adapter", meta["quantization"] |
| calibration = path / "adapter" / "calibration.json" |
| if not calibration.exists(): |
| calibration = path / "calibration.json" |
| model, tokenizer, letter_ids, _ = load_decision_model( |
| model_id, quantization=quantization or default_quantization, max_memory=max_memory, |
| adapter_dir=adapter_dir, counter=counter) |
| if letter_ids != head_config["letter_ids"]: |
| raise ValueError("this tokenizer's letter ids differ from the ones the head was trained with") |
| model.head.load_state_dict(load_file(str(path / "decision_head.safetensors"))) |
| model.head.to(next(model.backbone.parameters()).device) |
| temperatures = {kind: 1.0 for kind in QUESTION_TYPES} |
| if calibration.exists(): |
| temperatures.update(json.loads(calibration.read_text())["temperatures"]) |
| else: |
| counter["uncalibrated_temperature_1"] += 1 |
| return cls(model, tokenizer, temperatures, meta["max_length"], counter) |
|
|
| def decide(self, state: str, question: dict[str, Any], focus: str | None = None) -> dict[str, Any]: |
| """question: {"type": "choice"|"noul"|"score", "instructions": str, "criteria": {option_id: definition}}. |
| noul criteria are {"true": ..., "false": ...}; score criteria are {"0": lowest, "1": ..., ...}.""" |
| labels = ["true", "false"] if question["type"] == "noul" else list(question["criteria"]) |
| row = {"state": state, "question": question, "labels": labels, |
| "focus": {"surface": focus} if focus else {}} |
| logits, truncated = option_logits(self.model, self.tokenizer, row, self.max_length, self.fallback_counts) |
| probs = probabilities(logits, self.temperatures[question_type(row)]) |
| answer = jev_answer(row, probs) |
| answer["probabilities"] = probs |
| answer["state_truncated"] = truncated |
| return answer |
|
|
