File size: 6,370 Bytes
b6d5136 | 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 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 | """TeacherVariant — the interface every variant folder implements.
A variant is a self-contained folder under agents/variants/:
<variant>/utils.py a TeacherVariant subclass (+ VARIANT = TheSubclass)
<variant>/prompts.yaml its maker strategy / system (+ optional teacher wrapper)
<variant>/README.md what it is, its bet, its prediction
The SHARED machinery — HTTP transport (ChatClient.score_topk / complete), the MC
loop, the divergence metrics, and the batched training loop — is NOT here and is
never copied. The probe/trainer call the methods below, so adding a variant never
adds an `if variant == ...` branch to core code.
Each variant overrides ONLY what differs; the base is the current live behavior
(reference in the user turn, analyze inside a closed <think>, then derive).
"""
import os
import re
import difflib
import yaml
from agents.utils import (_fill, teacher_user_msg, ANALYSIS_LEADIN,
POST_THINK_TRANSITION)
from agents.maker import build_prefix
from agents.verify import is_acceptable_prefix, answer_matches
# Shared maker user-template — variants supply only strategy + system, and inherit
# this fixed scaffold (it is plumbing, not variant logic, so it is not duplicated).
_MAKER_TEMPLATE = yaml.safe_load(
open(os.path.join(os.path.dirname(__file__), "..", "prompts", "maker_prompts.yaml"))
)["maker_user_template"]
class TeacherVariant:
"""Base = the live 'thinking_on' behavior. Subclasses override selectively."""
name = "base"
description = "reference in the user turn; analyze in a closed <think>, then derive"
thinking = True # does the teacher analyze inside a <think> block?
uses_maker = True # does resolve_prefix call the maker endpoint? (P4: False)
def __init__(self, variant_dir=None):
if variant_dir is None:
import sys
mod = sys.modules.get(self.__class__.__module__)
variant_dir = os.path.dirname(getattr(mod, "__file__", "") or ".")
self.dir = variant_dir
self.prompts_path = os.path.join(self.dir, "prompts.yaml")
with open(self.prompts_path) as f:
self.prompts = yaml.safe_load(f) or {}
# ------------------------------------------------------------------ #
# 1. prefix resolution — what fills the reference slot of the context
# ------------------------------------------------------------------ #
def _maker_prompts(self):
return {"strategy": self.prompts["strategy"],
"maker_system_prompt": self.prompts["maker_system_prompt"],
"maker_user_template": _MAKER_TEMPLATE}
def resolve_prefix(self, problem, solution, rollout, maker_client):
"""Maker rewrite of the reference in the student's voice; falls back to the
reference on any failure. Uses THIS variant's strategy + accept policy."""
if maker_client is None or not (solution or "").strip():
return solution
return build_prefix(problem, solution, rollout, maker_client,
prompts=self._maker_prompts(), accept=self.accepts)
# ------------------------------------------------------------------ #
# 2. accept gate — is a candidate prefix good enough to use?
# ------------------------------------------------------------------ #
def accepts(self, prefix, reference):
return is_acceptable_prefix(prefix, reference)
# ------------------------------------------------------------------ #
# 3. teacher context (probe transport: needs the tokenizer + HTTP client)
# ------------------------------------------------------------------ #
def build_ctx(self, tok, client, problem, prefix, cfg):
wrapper = self.prompts.get("teacher_user")
user_msg = _fill(wrapper, problem, prefix) if wrapper else teacher_user_msg(problem, prefix)
base = tok.apply_chat_template(
[{"role": "user", "content": user_msg}], tokenize=False,
add_generation_prompt=True, enable_thinking=self.thinking)
if not self.thinking:
return base
lead = self.prompts.get("analysis_leadin", ANALYSIS_LEADIN)
gen_prompt = base + "<think>\n" + lead + "\n\n"
out = client.complete(gen_prompt, n=1, temperature=cfg["analysis_temperature"],
max_tokens=cfg["analysis_max_tokens"], stop=["</think>"])
analysis = (out[0] if out else "").split("</think>")[0].rstrip()
trans = self.prompts.get("post_think_transition", POST_THINK_TRANSITION)
return gen_prompt + analysis + "\n</think>\n\n" + trans
# ------------------------------------------------------------------ #
# 4. error localization — t* (token index into the rollout)
# ------------------------------------------------------------------ #
def t_star(self, prefix, rollout, tok):
return self._difflib_t_star(prefix, rollout, tok)
# -------- shared t* helpers (a variant picks one) -------- #
@staticmethod
def _char_to_tok(text, char_idx, tok):
if char_idx is None or char_idx < 0:
return None
return len(tok(text[:char_idx], add_special_tokens=False)["input_ids"])
@classmethod
def _difflib_t_star(cls, prefix, rollout, tok):
"""First block where the rollout stops matching the prefix (approx)."""
sm = difflib.SequenceMatcher(a=prefix or "", b=rollout or "", autojunk=False)
for tag, i1, i2, j1, j2 in sm.get_opcodes():
if tag != "equal":
return cls._char_to_tok(rollout, j1, tok)
return None
@classmethod
def _lcp_t_star(cls, prefix, rollout, tok):
"""Longest common prefix boundary — the splice point for verbatim-keep r."""
a, b = prefix or "", rollout or ""
i = 0
while i < min(len(a), len(b)) and a[i] == b[i]:
i += 1
return cls._char_to_tok(rollout, i, tok)
@classmethod
def _quote_t_star(cls, quote, rollout, tok):
"""Locate a quoted first-wrong-sentence inside the rollout (whitespace-tolerant)."""
if not quote:
return None
pat = re.compile(re.escape(quote.strip()).replace(r"\ ", r"\s+"))
m = pat.search(rollout or "")
return cls._char_to_tok(rollout, m.start(), tok) if m else None
|