281 lines
12 KiB
Python
281 lines
12 KiB
Python
"""LLM answering: prompts, generation, and output parsing (scaffolded and lean)."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import re
|
|
from typing import List, Optional, Sequence, Tuple
|
|
|
|
from .llm import LLMClient
|
|
from .preprocess import Puzzle
|
|
|
|
SYSTEM = (
|
|
"You solve International Linguistics Olympiad problems. Every problem is "
|
|
"about a language you have never seen; ALL the evidence you need is in the "
|
|
"given data. Derive the rules ONLY from that data — do not assume the "
|
|
"language works like English or any language you know. "
|
|
"A 'Mechanical analysis' section may provide segmentation, alignment, and "
|
|
"candidate answers computed by exact algorithms: use them as hypotheses, "
|
|
"adopt candidates that fit the data, and correct them when the data "
|
|
"disagrees. "
|
|
"First, reason briefly about the linguistic patterns — keep it to a few "
|
|
"lines, not an essay: line up the given examples, segment the words into "
|
|
"morphemes, note what each recurring morpheme means and the order they "
|
|
"combine in, and any sound changes; check that the pattern holds across the "
|
|
"examples, then apply it to each query item. "
|
|
"Answer format by task type -- translation: the translated form only, in "
|
|
"the language asked for; fill_blanks: only the missing form for each "
|
|
"blank; match_letters: only the option letter (for example A, B, C); "
|
|
"text_to_num: the number in digits; num_to_text: the number written out "
|
|
"in the puzzle language; anything else: exactly what the instruction "
|
|
"asks, nothing more. Match the punctuation style of the given examples. "
|
|
"Then write a line that says exactly "
|
|
"FINAL ANSWERS: and, below it, one answer per line in the order the items "
|
|
"are asked -- the bare answer only, no numbering, no quotes, no markdown, "
|
|
"no extra text. If you are unsure, still give your single best guess for "
|
|
"every item -- never leave an item blank. After the answers, write a line "
|
|
"that says exactly EXPLANATION: followed by 2-4 short bullet points "
|
|
"stating the rules you found (word order, affixes and their functions, "
|
|
"sound changes, numeral bases) and the key evidence for them. Do not "
|
|
"repeat your reasoning.\n\n"
|
|
"Output shape (follow it exactly; do not bold or number the markers):\n"
|
|
"<a few lines of pattern reasoning>\n"
|
|
"FINAL ANSWERS:\n"
|
|
"first answer\n"
|
|
"second answer\n"
|
|
"EXPLANATION:\n"
|
|
"- rule and evidence\n"
|
|
"- rule and evidence"
|
|
)
|
|
|
|
# Minimal prompt: no scaffold, no chain-of-thought. Selected by LEAN_MODE.
|
|
LEAN_SYSTEM = (
|
|
"You solve International Linguistics Olympiad problems. Answer every "
|
|
"numbered item. Put each answer on its own line, in order, with no "
|
|
"numbering and no extra text. Give your best guess for every item; never "
|
|
"leave one blank."
|
|
)
|
|
|
|
MAX_NEW_TOKENS = 1536
|
|
|
|
# Markers tolerate leading markdown/quote decoration (#, *, >, -, spaces) and
|
|
# trailing decoration/colon, and capture any inline content after the marker
|
|
# ("FINAL ANSWERS: foo" -> "foo" is the first answer). Qwen habitually bolds
|
|
# these headers; the strict "^marker$" form silently dropped every such output.
|
|
# trailing class excludes newline ([^\S\n]) so the marker never swallows the
|
|
# line break and misreads the next line as inline content
|
|
_FINAL_RX = re.compile(r"(?im)^[ \t#>*_`-]*final[ \t]*answers?\b[^\S\n]*[:*_`]?[^\S\n]*(.*)$")
|
|
_EXPL_RX = re.compile(r"(?im)^[ \t#>*_`-]*explanation\b[^\S\n]*[:*_`]?[^\S\n]*(.*)$")
|
|
_NUMBERING_RX = re.compile(r"^\s*\(?\d{1,3}[.)]\s*")
|
|
# leading markdown decoration on an answer line: bullets, bold, backticks
|
|
_ANSWER_DECOR_RX = re.compile(r"^[\s>*_`•·–-]+")
|
|
|
|
|
|
def build_prompt(puzzle: Puzzle, scaffold: str = "") -> str:
|
|
base = f"{puzzle.context.strip()}\n\n{puzzle.query.strip()}"
|
|
if scaffold:
|
|
base += f"\n\n{scaffold}"
|
|
return base
|
|
|
|
|
|
def _clean_answer_line(line: str) -> str:
|
|
line = _NUMBERING_RX.sub("", line.strip())
|
|
line = _ANSWER_DECOR_RX.sub("", line)
|
|
# strip trailing markdown emphasis but keep sentence punctuation (EM-significant)
|
|
line = re.sub(r"[*_`]+$", "", line)
|
|
return line.strip().strip("'\"“”").strip()
|
|
|
|
|
|
def parse_output(text: str, n_items: Optional[int] = None
|
|
) -> Tuple[List[str], str, bool]:
|
|
"""Returns (answers, explanation, found_marker).
|
|
|
|
`found_marker` is True only when an explicit FINAL ANSWERS marker was
|
|
present. On a parse failure (no marker) we return ([], expl_if_any, False)
|
|
rather than treating the reasoning prose as answers -- the caller then
|
|
keeps its symbolic answers instead of overwriting them with garbage.
|
|
When `n_items` is given it is used only as a sanity cap on how many
|
|
answer lines to accept (guards against a runaway list)."""
|
|
markers = list(_FINAL_RX.finditer(text))
|
|
found = bool(markers)
|
|
if markers:
|
|
m = markers[-1]
|
|
# a marker may carry its first answer inline: "FINAL ANSWERS: foo"
|
|
inline = m.group(1).strip() if m.lastindex else ""
|
|
tail = (inline + "\n" if inline else "") + text[m.end():]
|
|
else:
|
|
tail = ""
|
|
|
|
expl = ""
|
|
em = _EXPL_RX.search(tail) if found else None
|
|
if em:
|
|
inline_e = em.group(1).strip() if em.lastindex else ""
|
|
rest_e = tail[em.end():].strip()
|
|
expl = (inline_e + ("\n" + rest_e if rest_e else "")) if inline_e else rest_e
|
|
tail = tail[: em.start()]
|
|
|
|
answers: List[str] = []
|
|
cap = (2 * n_items + 4) if n_items else None
|
|
for line in tail.splitlines():
|
|
cleaned = _clean_answer_line(line)
|
|
if cleaned:
|
|
answers.append(cleaned)
|
|
if cap and len(answers) >= cap:
|
|
break
|
|
|
|
# models sometimes put EXPLANATION before FINAL ANSWERS despite the prompt
|
|
if not expl:
|
|
em2 = _EXPL_RX.search(text)
|
|
if em2 and (not markers or em2.start() < markers[-1].start()):
|
|
seg = text[em2.end():]
|
|
stop = _FINAL_RX.search(seg)
|
|
expl = (seg[: stop.start()] if stop else seg).strip()
|
|
return answers, expl, found
|
|
|
|
|
|
# lines the model may prepend/append around a bare answer list
|
|
_CHATTY_RX = re.compile(
|
|
r"^(?:here (?:are|is)\b|answers?\s*:?\s*$|the answers?\b|translations?\s*:?\s*$|"
|
|
r"note\b|okay\b|sure\b|solution\b|let me\b)", re.IGNORECASE)
|
|
|
|
|
|
def parse_output_lean(text: str, n_items: Optional[int] = None
|
|
) -> Tuple[List[str], str, bool]:
|
|
"""Lean parse: the minimal prompt asks for bare answers, one per line, with
|
|
no FINAL ANSWERS marker — so every non-empty line IS an answer. We only
|
|
strip numbering/markdown decoration and drop obvious chatty preamble lines.
|
|
When more than n_items lines survive, keep the LAST n_items (any stray
|
|
preamble sits at the top). No marker requirement, no salvage."""
|
|
if not text:
|
|
return [], "", False
|
|
answers: List[str] = []
|
|
for line in text.splitlines():
|
|
cleaned = _clean_answer_line(line)
|
|
if cleaned and not _CHATTY_RX.match(cleaned):
|
|
answers.append(cleaned)
|
|
if n_items and len(answers) > n_items:
|
|
answers = answers[-n_items:]
|
|
return answers, "", bool(answers)
|
|
|
|
|
|
_LEAN_LABEL_RX = re.compile(r"^\s*\(?(\d{1,3})\)?[.):\]]\s")
|
|
|
|
|
|
def align_lean(text: str, puzzle: Puzzle) -> List[Optional[str]]:
|
|
"""Turn a lean (bare-lines) model output into exactly len(puzzle.items)
|
|
answers, blending two placement methods that stack:
|
|
|
|
1. LABEL-AWARE placement — if the model numbered its answer lines and those
|
|
numbers cover most of the item labels, place each answer under its own
|
|
label. This is robust to the model reordering items or skipping one (a
|
|
single skipped line would otherwise shift every later answer and zero the
|
|
whole block on both metrics).
|
|
2. POSITIONAL last-N fallback — when the output isn't reliably numbered, take
|
|
the cleaned, non-chatty lines in order (dropping any preamble at the top).
|
|
|
|
Both share our line hygiene (numbering/markdown stripping via
|
|
`_clean_answer_line`, preamble drop via `_CHATTY_RX`). Returns a length-N
|
|
list; None marks items the model did not answer, which the caller leaves to
|
|
the fallback."""
|
|
items = puzzle.items
|
|
n = len(items)
|
|
if n == 0:
|
|
return []
|
|
labeled: dict = {}
|
|
for line in (text or "").splitlines():
|
|
s = line.strip()
|
|
if not s:
|
|
continue
|
|
m = _LEAN_LABEL_RX.match(s)
|
|
cleaned = _clean_answer_line(s)
|
|
if not cleaned or _CHATTY_RX.match(cleaned):
|
|
continue
|
|
if m:
|
|
labeled[m.group(1)] = cleaned # last write wins (models restate)
|
|
item_labels = [it.number for it in items]
|
|
if labeled and all(lbl for lbl in item_labels):
|
|
covered = sum(1 for lbl in item_labels if lbl in labeled)
|
|
if covered >= max(1, (2 * n + 2) // 3): # ~2/3 labelled -> trust labels
|
|
return [labeled.get(lbl) for lbl in item_labels]
|
|
# positional fallback: our bare-line parse, aligned by position
|
|
bare, _e, _f = parse_output_lean(text, n)
|
|
out: List[Optional[str]] = [None] * n
|
|
for i in range(min(n, len(bare))):
|
|
out[i] = bare[i]
|
|
return out
|
|
|
|
|
|
def build_salvage_prompt(base_prompt: str, raw_reasoning: str,
|
|
max_chars: int = 2400) -> str:
|
|
"""A short follow-up prompt that reuses reasoning the model already
|
|
produced (but never terminated with a FINAL ANSWERS block, e.g. it hit the
|
|
token cap). We feed the reasoning back as context and ask ONLY for the
|
|
answer block -- cheaper and more reliable than a bare retry."""
|
|
reasoning = raw_reasoning.strip()
|
|
if len(reasoning) > max_chars:
|
|
reasoning = reasoning[-max_chars:] # keep the most recent reasoning
|
|
return (
|
|
f"{base_prompt}\n\n"
|
|
"You already worked through this problem:\n"
|
|
"-----\n"
|
|
f"{reasoning}\n"
|
|
"-----\n"
|
|
"Now output ONLY the answer block, nothing else. Write the line "
|
|
"FINAL ANSWERS: then one bare answer per line in item order (your best "
|
|
"guess for every item, never blank), then a line EXPLANATION: with 2-4 "
|
|
"short bullets."
|
|
)
|
|
|
|
|
|
def solve_direct(puzzles: Sequence[Puzzle], client: LLMClient,
|
|
scaffolds: Optional[Sequence[str]] = None,
|
|
max_new_tokens: int = MAX_NEW_TOKENS,
|
|
salvage: bool = True,
|
|
system: str = SYSTEM,
|
|
lean: bool = False,
|
|
sample: bool = False,
|
|
temperature: float = 0.5
|
|
) -> List[Tuple[List[str], str, str, bool]]:
|
|
"""Batched single-shot answering. Returns per puzzle
|
|
(answers, explanation, raw_text, found_marker). `sample`/`temperature` drive
|
|
the sampled passes used by self-consistency voting; salvage is skipped when
|
|
sampling. In `lean` mode every non-empty line is an answer (no marker)."""
|
|
if not client.available or not puzzles:
|
|
return [([], "", "", False) for _ in puzzles]
|
|
prompts = [
|
|
build_prompt(p, scaffolds[i] if scaffolds else "")
|
|
for i, p in enumerate(puzzles)
|
|
]
|
|
raws = client.generate(prompts, max_new_tokens=max_new_tokens, system=system,
|
|
sample=sample, temperature=temperature)
|
|
out: List[Tuple[List[str], str, str, bool]] = []
|
|
salvage_idx: List[int] = []
|
|
for i, r in enumerate(raws):
|
|
if lean:
|
|
aligned = align_lean(r, puzzles[i])
|
|
answers, expl, found = aligned, "", any(a for a in aligned)
|
|
else:
|
|
answers, expl, found = parse_output(r, len(puzzles[i].items))
|
|
out.append((answers, expl, r, found))
|
|
if salvage and not lean and not sample and not found and r.strip():
|
|
salvage_idx.append(i)
|
|
|
|
if salvage_idx:
|
|
sp = [build_salvage_prompt(prompts[i], raws[i]) for i in salvage_idx]
|
|
# answers only — a small cap keeps the salvage pass cheap
|
|
cap = min(max_new_tokens, 384)
|
|
sraws = client.generate(sp, max_new_tokens=cap, system=system)
|
|
for i, sr in zip(salvage_idx, sraws):
|
|
a2, e2, f2 = parse_output(sr, len(puzzles[i].items))
|
|
if f2 and a2:
|
|
prev = out[i]
|
|
out[i] = (a2, e2 or prev[1], prev[2], True)
|
|
return out
|
|
|
|
|
|
def align_answers(direct: List[str], n_items: int) -> List[Optional[str]]:
|
|
"""Position-align a FINAL ANSWERS block to the expected item count."""
|
|
out: List[Optional[str]] = [None] * n_items
|
|
for i in range(min(n_items, len(direct))):
|
|
out[i] = direct[i]
|
|
return out
|