File size: 3,746 Bytes
9ebcba2 | 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 84 85 86 87 88 89 | #!/usr/bin/env python3
"""Yes/no forecast. One forward pass, then the fitted map. Returns P(YES) only."""
from __future__ import annotations
import math
import string
SYSTEM_PROMPT = "Apply the criterion to the evidence. Choose exactly one listed option. Answer with its letter only."
THINK_OFF = "<think>\n\n</think>\n\n"
TEMPERATURE = 1.421
# Fitted on the primary-election evaluation.
A = 1.099406044975045
C = 1.3783573144958756
def _logit(p: float) -> float:
p = min(max(p, 1e-6), 1 - 1e-6)
return math.log(p / (1 - p))
def _sigmoid(x: float) -> float:
return 1.0 / (1.0 + math.exp(-x))
def _letter_ids(tokenizer) -> list[int]:
probe = [
{"role": "system", "content": SYSTEM_PROMPT},
{"role": "user", "content": "Evidence:\n(none)\n\nQuestion: q\nOptions:\nA) a\nB) b"},
]
text = tokenizer.apply_chat_template(probe, tokenize=False, add_generation_prompt=True, enable_thinking=False)
if not text.endswith(THINK_OFF):
raise ValueError("chat template is missing the expected end marker")
base = tokenizer.encode(text, add_special_tokens=False)
ids = []
for letter in string.ascii_uppercase:
piece = tokenizer.encode(letter, add_special_tokens=False)
if len(piece) != 1 or tokenizer.encode(text + letter, add_special_tokens=False) != base + piece:
raise ValueError(f"{letter!r} is not a single token for this tokenizer")
ids.append(piece[0])
return ids
class Forecaster:
"""Load once. forecast(...) returns a float in [0, 1]."""
def __init__(self, model_path: str = "YarrowLab/yarrow-forecast"):
import torch
from transformers import AutoTokenizer, Qwen3_5ForConditionalGeneration
self._torch = torch
self.tok = AutoTokenizer.from_pretrained(model_path)
self.model = Qwen3_5ForConditionalGeneration.from_pretrained(
model_path, dtype="auto", device_map="auto",
).eval()
self._yes, self._no = _letter_ids(self.tok)[:2]
def forecast(self, title: str, rules: str, evidence: str = "", as_of: str | None = None) -> float:
return self.forecast_batch([{
"title": title, "rules": rules, "evidence": evidence, "as_of": as_of,
}])[0]
def forecast_batch(self, questions: list[dict]) -> list[float]:
return [self._one(q) for q in questions]
def _one(self, q: dict) -> float:
as_of = q.get("as_of") or "the question date"
state = f"Forecast as of {as_of}."
if q.get("evidence"):
state += "\n\nEVIDENCE AVAILABLE AT FORECAST TIME:\n" + q["evidence"]
question = (
f"{q['title']}\n\nResolution rules (verbatim): {q['rules']}\n\n"
"Will this resolve YES, based only on the evidence above?"
)
state = state if state.strip() else "(none)"
messages = [
{"role": "system", "content": SYSTEM_PROMPT},
{"role": "user", "content": "Evidence:\n" + state + "\n\nQuestion: " + question.strip() + "\nOptions:\nA) yes\nB) no"},
]
text = self.tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, enable_thinking=False)
if not text.endswith(THINK_OFF):
raise ValueError("chat template is missing the expected end marker")
ids = self.tok.encode(text, add_special_tokens=False)
input_ids = self._torch.tensor([ids], device=self.model.device)
with self._torch.no_grad():
logits = self.model(input_ids=input_ids).logits[0, -1].float()
pair = self._torch.stack([logits[self._yes], logits[self._no]])
yes = self._torch.softmax(pair / TEMPERATURE, dim=0)[0].item()
return _sigmoid(A * _logit(yes) + C)
|