| |
| """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 |
| |
| 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) |
|
|
