File size: 4,496 Bytes
53ab3d4 cb5578c 53ab3d4 cb5578c aa20e73 53ab3d4 aa20e73 53ab3d4 aa20e73 53ab3d4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 | """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:
# Full model: the LoRA is merged into the weights in this folder; no base download, no peft.
model_id, adapter_dir, default_quantization = str(path), None, "none"
calibration = path / "calibration.json"
else:
# Adapter on the base model, 4-bit as trained.
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
|