"""Pipeline (shared by script.py and eval/dev_harness.py): symbolic pass, then LLM answering (lean or scaffolded), merge, never-empty guarantee.""" from __future__ import annotations import re from dataclasses import dataclass, field from typing import Callable, List, Optional, Sequence from .budget import Budget from .direct import align_answers, solve_direct from .llm import LLMClient, NullClient from .matching import solve_match_letters_llm from .preprocess import Puzzle, parse_puzzle from .router import solve_puzzle_ex from .scaffold import build_scaffold, light_hint CONF_KEEP = 0.5 # symbolic answers at/above this verifier fit override the LLM LLM_BATCH = 4 # puzzles per generation batch (T4 KV-cache friendly) DEEP_MODE_MAX_ROWS = 30 # at/below this many puzzles, spend more per puzzle _DIGITS_RX = re.compile(r"\d[\d,. ]*") _LETTER_RX = re.compile(r"\b([A-Z])\b") _PREAMBLE_RX = re.compile( r"^(?:the\s+)?(?:answer|translation|result)\s*(?:is|:)\s*", re.IGNORECASE) _TERMINAL_PUNCT = (".", "!", "?") @dataclass class PuzzleResult: row_id: str answers: List[str] explanation: str = "" confs: List[float] = field(default_factory=list) methods: List[str] = field(default_factory=list) llm_used: bool = False raw: str = "" def clean_llm_answer(ans: str, task_type: str, option_labels: Sequence[str] = ()) -> str: """Per-task-type format guard on a parsed LLM answer line. Conservative: only rewrites when the expected shape is unambiguous.""" a = _PREAMBLE_RX.sub("", ans.strip()).strip().strip("'\"“”") if task_type == "text_to_num": m = _DIGITS_RX.search(a) if m: digits = re.sub(r"[,. ]", "", m.group(0)) if digits.isdigit(): return digits elif task_type == "match_letters": if len(a) > 2: # "B. the bird sleeps" or "option B" -> "B" candidates = _LETTER_RX.findall(a) wanted = [c for c in candidates if not option_labels or c in option_labels] if len(set(wanted)) == 1: return wanted[0] return a def _vote(cands: Sequence[Optional[str]], anchor: Optional[str]) -> Optional[str]: """Greedy-anchored vote: keep `anchor` (the greedy answer) unless at least two sampled candidates agree on the same normalised form AND that form outnumbers the anchor's support. Monotone — it can only fire on genuine agreement, so it never replaces greedy with a lone sample.""" from collections import Counter cands = [c for c in cands if c and str(c).strip()] if anchor is None: anchor = cands[0] if cands else None if len(cands) < 3: return anchor norm = lambda s: re.sub(r"\s+", " ", str(s).strip().lower()) groups: dict = {} for c in cands: groups.setdefault(norm(c), []).append(c) anchor_support = len(groups.get(norm(anchor), [])) if anchor else 0 best_key = max(groups, key=lambda k: len(groups[k])) if len(groups[best_key]) >= 2 and len(groups[best_key]) > anchor_support: return Counter(groups[best_key]).most_common(1)[0][0] return anchor def induce_format(answer: str, puzzle: Puzzle, item_idx: int) -> str: """Nudge an answer toward the dataset's surface convention (fable §3): if the attested answers on this item's side overwhelmingly end in a terminal punctuation mark or start with a capital, mirror that. Converts chrF-close answers into EM hits, which the geometric-mean scoring rewards twice. Conservative: only fires on a near-unanimous (>=85%) convention, only ADDS a missing terminal mark or leading capital, never strips.""" if puzzle.task_type not in ("translation", "fill_blanks"): return answer a = answer.strip() if not a: return answer it = puzzle.items[item_idx] if item_idx < len(puzzle.items) else None direction = getattr(it, "direction", None) # answer side: to_work -> work-language (tgt); else task-language (src) side = [p.tgt for p in puzzle.pairs] if direction == "to_work" \ else [p.src for p in puzzle.pairs] side = [s.strip() for s in side if s and s.strip()] if len(side) < 4: return answer n = len(side) # terminal punctuation: only if a single mark dominates for mark in _TERMINAL_PUNCT: if sum(1 for s in side if s.endswith(mark)) / n >= 0.85: if not a.endswith(_TERMINAL_PUNCT): a = a + mark break # leading capitalization if sum(1 for s in side if s[:1].isupper()) / n >= 0.85: if a[:1].islower(): a = a[:1].upper() + a[1:] return a def _symbolic_explanation(puzzle: Puzzle, methods: Sequence[str]) -> str: used = [m for m in dict.fromkeys(methods) if m not in ("none", "fallback")] if not used: return ("- Answered by nearest-attested analogy over the given examples " "(no reliable rule could be verified).") tmpl = { "numeral system (verified)": ( "- Induced each morpheme's numeric value from the attested numerals " "and verified the system reproduces every given example; applied it " "to each query item (smaller-before-larger multiplies, otherwise " "values add)."), "table completion": ( "- Learned the mapping between the paradigm-table columns from the " "attested rows (leave-one-out verified) and applied it to each " "incomplete row."), "template substitution": ( "- For each query, took the closest attested sentence and swapped " "the differing words through morpheme alignments induced from " "minimal pairs in the data."), "induced grammar": ( "- Induced a lexicon and affix rules that reproduce the attested " "pairs exactly, then applied them mechanically to the query items."), "optimal matching": ( "- Scored every form-meaning pair by shared-morpheme/shared-word " "consistency and picked the globally optimal assignment."), } return "\n".join(tmpl.get(m, f"- Solved by {m}.") for m in used) def run_pipeline(rows: Sequence[dict], client: Optional[LLMClient] = None, budget: Optional[Budget] = None, verbose: bool = True, conf_keep: float = CONF_KEEP, llm_batch: int = LLM_BATCH, max_new_tokens: Optional[int] = None, checkpoint: Optional[Callable[[List["PuzzleResult"]], None]] = None, lean: bool = False, use_match_assignment: bool = True, vote_samples: int = 0, vote_temp: float = 0.5, hint: bool = False ) -> List[PuzzleResult]: """`checkpoint`, when given, is called with the (complete, valid) results after the symbolic pass and after every LLM batch — so a crash at ANY later point still leaves a full submission on disk.""" client = client or NullClient() budget = budget or Budget() # adaptive mode: the hidden test is small (an IOL contest reformatted into # a handful of multi-item rows), so default to spending more per puzzle. # A large row count flips us to coverage mode (shorter generations, serve # the most puzzles). deep_mode = len(rows) <= DEEP_MODE_MAX_ROWS if max_new_tokens is None: max_new_tokens = 2048 if deep_mode else 1024 def log(msg: str) -> None: if verbose: print(msg, flush=True) def save() -> None: if checkpoint is not None: try: checkpoint(results) except Exception: pass # ---- 1. symbolic pass ---- results: List[PuzzleResult] = [] puzzles: List[Optional[Puzzle]] = [] for i, row in enumerate(rows): rid = str(row.get("id", i)) try: p = parse_puzzle(row) answers, confs, methods = solve_puzzle_ex(p, NullClient(), budget, puzzles_left=len(rows) - i) except Exception: p = None answers, confs, methods = [str(row.get("query", "?")).strip() or "?"], [0.0], ["fallback"] puzzles.append(p) results.append(PuzzleResult(rid, answers, "", confs, methods)) # make every result submission-valid NOW (explanations + non-empty), so # each checkpoint from here on is a complete fallback submission for i, r in enumerate(results): r.answers = [str(a).strip() or "?" for a in r.answers] r.explanation = (_symbolic_explanation(puzzles[i], r.methods) if puzzles[i] is not None else "- No parseable structure found; answered by " "closest-example analogy.") log(f"symbolic pass done in {budget.elapsed():.1f}s") save() # arm the per-token wall-clock abort so generation can't overrun the budget try: client.deadline = budget.start + budget.total - budget.safety except Exception: pass # ---- 1b. match_letters assignment pass (model logprobs -> Hungarian) ---- # Free-form generation answers match_letters with the identity permutation # (scores ~0). Solve it as an assignment from the model's own distribution # instead. Solved puzzles get high confidence so the free-form LLM pass # skips them; a declined puzzle falls through to that pass unchanged. # Gated by use_match_assignment: when off, match_letters puzzles go through # the normal free-form LLM pass. if not use_match_assignment: log("match_letters assignment pass disabled; using free-form LLM path") if use_match_assignment and getattr(client, "can_score", False): n_assigned = 0 for i, p in enumerate(puzzles): if p is None or p.task_type != "match_letters" or budget.exhausted(): continue try: letters = solve_match_letters_llm(p, client) except Exception as e: log(f" match_letters solver failed on {results[i].row_id}: " f"{type(e).__name__}: {e}") letters = None r = results[i] if letters and len(letters) == len(r.answers): r.answers = [str(x).strip() or "?" for x in letters] r.confs = [0.9] * len(letters) r.methods = ["llm-assignment"] * len(letters) r.llm_used = True n_assigned += 1 if n_assigned: log(f"match_letters assignment solver used on {n_assigned} puzzle(s)") save() # ---- 2. scaffolded LLM pass ---- # In lean mode symbolic is a pure fallback: the model answers EVERY puzzle # and its answer wins wherever it produced one (symbolic stands only for the # items it left blank). Otherwise the model runs only on low-confidence # puzzles and verified symbolic answers override it. if lean: need = [i for i in range(len(results)) if puzzles[i] is not None] else: need = [i for i, r in enumerate(results) if puzzles[i] is not None and any(c < conf_keep for c in r.confs)] if deep_mode: # weakest-first: a budget cutoff then drops the puzzles we could help least need.sort(key=lambda i: sum(min(c, conf_keep) for c in results[i].confs) / max(len(results[i].confs), 1)) else: # coverage: shortest prompt first maximizes puzzles served per second need.sort(key=lambda i: len(puzzles[i].context) + len(puzzles[i].query)) from .direct import LEAN_SYSTEM, SYSTEM sys_prompt = LEAN_SYSTEM if lean else SYSTEM log(f"LLM pass ({'lean' if lean else 'scaffold'}, " f"{'deep' if deep_mode else 'coverage'}, " f"max_new_tokens={max_new_tokens}): {len(need)}/{len(rows)} puzzles " f"need the model; client={'yes' if client.available else 'no'}") parse_ok = parse_fail = 0 if client.available and need: done = 0 for start in range(0, len(need), llm_batch): if budget.exhausted(): log(f"budget cutoff after {done} LLM puzzles") break batch = need[start : start + llm_batch] if lean: # lean: no scaffold, unless the optional light hint is enabled scaffolds = [light_hint(puzzles[i]) if hint else "" for i in batch] else: scaffolds = [] for i in batch: p, r = puzzles[i], results[i] try: scaffolds.append(build_scaffold(p, r.answers, r.confs, r.methods)) except Exception: scaffolds.append("") try: outs = solve_direct([puzzles[i] for i in batch], client, scaffolds, max_new_tokens=max_new_tokens, system=sys_prompt, lean=lean) except Exception as e: # systemic generation failure (bad load, driver, etc.) — the # symbolic answers already on every item are the submission log(f"LLM batch failed ({type(e).__name__}: {e}); " f"keeping symbolic answers for the rest") break for i, (direct, expl, raw, found) in zip(batch, outs): r = results[i] p = puzzles[i] r.llm_used = True r.raw = raw if found: parse_ok += 1 else: # no answer block parsed even after salvage: do NOT let # reasoning prose overwrite the symbolic answers parse_fail += 1 continue labels = sorted(p.lettered) if p.lettered else () aligned = align_answers(direct, len(r.answers)) for j, d in enumerate(aligned): # lean: the model's answer wins wherever it gave one; # scaffold: only where symbolic isn't confident if d and (lean or r.confs[j] < conf_keep): cleaned = clean_llm_answer(d, p.task_type, labels) if cleaned: # lean mode ships the model's answer as-is (no # punctuation/casing induction) if not lean: cleaned = induce_format(cleaned, p, j) r.answers[j] = cleaned.strip() or r.answers[j] r.methods[j] = "llm" if expl: r.explanation = expl done += len(batch) log(f" llm {done}/{len(need)} ok={parse_ok} fail={parse_fail} " f"t={budget.elapsed():.0f}s") save() log(f"LLM pass done: {parse_ok} parsed, {parse_fail} unparsable " f"(kept symbolic); {budget.elapsed():.0f}s elapsed") # ---- 2b. light greedy-anchored self-consistency voting (lean only) ---- # The greedy answers are already checkpointed; sampled passes can only # displace an item on genuine agreement (see _vote), so this is monotone and # budget-gated — if the clock runs out we simply keep the greedy answers. if lean and vote_samples > 0 and client.available and need and not budget.exhausted(): ballots = {i: [list(results[i].answers)] for i in need} # greedy = ballot 0 done_votes = 0 for _s in range(vote_samples): if budget.exhausted(): break failed = False for start in range(0, len(need), llm_batch): if budget.exhausted(): break batch = need[start : start + llm_batch] try: outs = solve_direct([puzzles[i] for i in batch], client, ["" for _ in batch], max_new_tokens=max_new_tokens, system=sys_prompt, lean=True, sample=True, temperature=vote_temp) except Exception as e: log(f"vote pass failed ({type(e).__name__}: {e}); " f"keeping greedy answers") failed = True break for i, (direct, _e, _raw, _f) in zip(batch, outs): p = puzzles[i] labels = sorted(p.lettered) if p.lettered else () ballot = [] for j in range(len(results[i].answers)): d = direct[j] if j < len(direct) else None ballot.append(clean_llm_answer(d, p.task_type, labels) if d else None) ballots[i].append(ballot) if failed: break done_votes += 1 for i in need: # re-vote (greedy-anchored) after each sample pass r = results[i] anchor = ballots[i][0] for j in range(len(r.answers)): v = _vote([b[j] for b in ballots[i] if j < len(b)], anchor[j] if j < len(anchor) else None) if v and str(v).strip(): r.answers[j] = str(v).strip() save() log(f"voting: {done_votes}/{vote_samples} sample pass(es); " f"t={budget.elapsed():.0f}s") # ---- 3. final never-empty guarantee ---- for r in results: r.answers = [str(a).strip() or "?" for a in r.answers] return results