linguist / script.py

Upload script.py

54b900b verified 4 months ago

15.5 kB

"""IOL-AI 2026 submission script.
Runs fully offline against /tmp/data/test.csv and writes submission.csv with
columns: id, pred (JSON list of answer strings), explanation (omitted).
"""
import os

os.environ["HF_HUB_OFFLINE"] = "1"
os.environ["TRANSFORMERS_OFFLINE"] = "1"

import time

T_START = time.monotonic()

import json
import re
from collections import Counter

import pandas as pd
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# --- Configuration -----------------------------------------------------

# Production always loads from "." (weights shipped alongside this script in
# the HF model repo). This constant only documents which Hub model those
# weights came from, for keeping the Colab test notebook in sync.
MODEL_ID_FOR_LOCAL_TEST = "Qwen/Qwen2.5-7B-Instruct-AWQ"
# Fallback if Colab timing shows 7B has slack to spare: swap MODEL_ID_FOR_LOCAL_TEST
# to "Qwen/Qwen2.5-14B-Instruct-AWQ" and re-run the notebook before re-uploading.

MODEL_PATH = "."
DATA_PATH = "/tmp/data/test.csv"
OUTPUT_PATH = "submission.csv"

TIME_BUDGET_SEC = 27 * 60  # 3-min safety margin under the 30-min hard limit
BATCH_SIZE = 4
MIN_NEW_TOKENS = 128
MAX_NEW_TOKENS = 1024

# Self-consistency: sample k >= 1 completions per row and majority-vote per
# answer position when the time budget has slack. k=1 is plain greedy decode
# (current default, deterministic). See adaptive_k().
MIN_K = 1
MAX_K = 1
SAMPLE_TEMPERATURE = 0.7

# Which SYSTEM prompt variant to use — "full_cot" or "plain_io". Published
# evidence on whether heavy chain-of-thought helps or hurts a 7B-class model
# on this task type is mixed (modeLing, arXiv:2406.17038), so both variants
# are kept and should be A/B tested on the Colab harness rather than assumed.
PROMPT_VARIANT = "plain_io"

# Worked example shared by both prompt variants: a made-up toy language (not
# a real one) demonstrating (a) the answer must be a word-form in the target
# language, never an English gloss, and (b) subject/object roles must be
# re-derived per query, not copied from a similar-looking example. Directly
# targets the two failure modes observed in real Colab runs (language
# confusion — arXiv:2406.20052 shows few-shot examples largely eliminate it;
# and swapped argument roles).
_WORKED_EXAMPLE = """Worked example (a made-up toy language, unrelated to your actual task — the
words and rules below apply ONLY to this example; never reuse them):
CONTEXT:
mota = "dog"        mota-ta = "dog" (as object)
suno = "cat"         suno-ta = "cat" (as object)
kire = "sees"
mota suno-ta kire = "the dog sees the cat"
suno mota-ta kire = "the cat sees the dog"
QUERY:
Translate into the made-up language: "the cat sees the dog"
IDENTIFY PATTERNS: word order is SUBJECT OBJECT-ta VERB; subject is unmarked,
object takes suffix "-ta", verb is always last.
RULE TABLE:
- subject noun -> bare noun, placed first
- object noun -> noun + "-ta", placed second
- verb -> placed last, unchanged
TEST: "mota suno-ta kire" = dog(subj) cat-ta(obj) sees = "the dog sees the
cat" — matches. "suno mota-ta kire" = cat(subj) dog-ta(obj) sees = "the cat
sees the dog" — matches.
APPLY: query is "the cat sees the dog" -> subject=cat=suno,
object=dog=mota-ta, verb=kire.
VERIFY: answer is a word-form in the made-up language, not an English gloss.
Roles checked: cat is subject (first, unmarked), dog is object (second,
"-ta" suffix) — matches "the cat sees the dog", not reversed.
FINAL ANSWERS:
suno mota-ta kire
--- end worked example. Now solve the actual task below using ONLY the data
in its own CONTEXT and QUERY — never reuse the words or rules above. ---"""

_ANSWER_FORMAT_RULES = """Answer formatting rules by task_type:
- translation: give the full translated phrase or sentence for each numbered item,
  written as a natural fluent sentence exactly like the answer column in the given
  examples — never as a morpheme-by-morpheme gloss with parentheses like
  "you(shuddered)" or "we(spat(on him))".
- fill_blanks: give only the missing word(s)/form for each numbered blank.
- match_letters: give the matching letter or number for each item.
- text_to_num: give only the numeral for each item.
- num_to_text: give only the word(s) for each item.
Output a line that says exactly:
FINAL ANSWERS:
followed by one answer per line, in the same order and count as the numbered
items in the query, with no numbering, quotes, or extra commentary — just the
bare answer text for each line."""

SYSTEM_FULL_COT = f"""You are an expert linguist solving International Linguistics Olympiad problems.
You are given a self-contained set of data in a language you have never seen before,
plus a query asking you to apply what that data teaches you. You must reason only from
the data given — never from memorized knowledge of real-world languages.
Work in stages, showing your work:
1. IDENTIFY PATTERNS: look for recurring forms, affixes, word order, sound
   correspondences, or structural regularities in the given data.
2. RULE TABLE: write your hypotheses as short "TRIGGER -> TRANSFORMATION"
   bullet lines (e.g. "subject noun -> bare noun, placed first"), not prose —
   this keeps rules mechanical and easy to re-check, rather than vague
   descriptions that drift away from the actual data.
3. TEST HYPOTHESES: check each rule-table line against every example in the
   given data. Discard or refine any rule that doesn't hold up on every
   example.
4. APPLY: use only the rules that survived testing to answer the query.
5. VERIFY: before writing FINAL ANSWERS, re-check every draft answer against
   two common mistakes:
   - LANGUAGE: look at the answer column in the given examples, not the
     instructions, to see what language/form your answer must be in. If the
     examples' answers are word-forms in the unfamiliar language, your answer
     must also be a word-form in that language — never substitute an English
     gloss or description of the meaning, even if you're unsure of the exact
     form; give your best-guess constructed form instead.
   - ROLES: for anything involving "who did what to whom" (subject, object,
     possessor, giver/receiver), re-read the query and re-derive which
     argument fills which role from scratch. Do not assume the same word
     order or role assignment as a similar-looking training example — verify
     it against the actual affixes/markers in that example.
{_WORKED_EXAMPLE}
{_ANSWER_FORMAT_RULES}"""

SYSTEM_PLAIN_IO = f"""You are an expert linguist solving International Linguistics Olympiad problems.
You are given a self-contained set of data in a language you have never seen before,
plus a query asking you to apply what that data teaches you. Reason only from the
data given — never from memorized knowledge of real-world languages. Give your
best-guess answer in the same language/form as the examples' answers — never
substitute an English gloss. Double-check which argument is subject vs object
before answering.
{_WORKED_EXAMPLE}
{_ANSWER_FORMAT_RULES}"""

_SYSTEM_VARIANTS = {"full_cot": SYSTEM_FULL_COT, "plain_io": SYSTEM_PLAIN_IO}
SYSTEM = _SYSTEM_VARIANTS[PROMPT_VARIANT]

USER_TEMPLATE = """Task type: {task_type}
Eval type: {eval_type}
Working language -> Task language: {work_lang} -> {task_lang}
CONTEXT:
{context}
QUERY:
{query}"""


# --- Query parsing helpers ---------------------------------------------

_RANGE_RE = re.compile(r"\((\d+)\s*[-–]\s*(\d+)\)")
_LEADING_NUM_RE = re.compile(r"^\s*(\d+)[.)]", re.MULTILINE)
_INLINE_PAREN_NUM_RE = re.compile(r"\((\d+)\)")


def _is_contiguous(nums: list[int]) -> bool:
    if not nums:
        return False
    uniq = sorted(set(nums))
    return uniq == list(range(uniq[0], uniq[-1] + 1))


def count_expected_items(query: str) -> int:
    """Count how many answers the query expects, robust to numbering style."""
    range_match = _RANGE_RE.search(query)
    if range_match:
        start, end = int(range_match.group(1)), int(range_match.group(2))
        if end >= start:
            return end - start + 1

    leading_nums = [int(n) for n in _LEADING_NUM_RE.findall(query)]
    if _is_contiguous(leading_nums):
        return len(set(leading_nums))

    inline_nums = [int(n) for n in _INLINE_PAREN_NUM_RE.findall(query)]
    if _is_contiguous(inline_nums):
        return len(set(inline_nums))

    non_empty_lines = [ln for ln in query.splitlines() if ln.strip()]
    return max(len(non_empty_lines) - 1, 1)


# --- Answer parsing helpers ----------------------------------------------

_MARKER_RE = re.compile(r"FINAL ANSWERS:\s*", re.IGNORECASE)
_LINE_NUM_PREFIX_RE = re.compile(r"^\s*\(?\d+\)?[.):]?\s*")
_QUOTE_STRIP_RE = re.compile(r'^["\'](.*)["\']$')


def _strip_line(line: str) -> str:
    line = line.strip()
    line = _LINE_NUM_PREFIX_RE.sub("", line, count=1)
    m = _QUOTE_STRIP_RE.match(line)
    if m:
        line = m.group(1)
    return line.strip()


def parse_answers(text: str, expected_n: int) -> list[str]:
    """Extract exactly expected_n answer strings from raw model output."""
    parts = _MARKER_RE.split(text, maxsplit=1)
    tail = parts[1] if len(parts) > 1 else text

    lines = [_strip_line(ln) for ln in tail.splitlines()]
    answers = [ln for ln in lines if ln]

    # Model sometimes emits all answers on one "|"-delimited line instead of
    # one per line; only fall back to splitting on "|" when the plain
    # line-split came up short, so genuine single-line answers containing "|"
    # aren't mangled.
    if len(answers) < expected_n and any("|" in a for a in answers):
        expanded = []
        for a in answers:
            if "|" in a:
                expanded.extend(_strip_line(p) for p in a.split("|"))
            else:
                expanded.append(a)
        expanded = [a for a in expanded if a]
        if len(expanded) > len(answers):
            answers = expanded

    if len(answers) < expected_n:
        answers = answers + [""] * (expected_n - len(answers))
    elif len(answers) > expected_n:
        answers = answers[:expected_n]

    return answers


# --- Time-budget-aware batched generation --------------------------------

def elapsed() -> float:
    return time.monotonic() - T_START


def remaining_budget() -> float:
    return TIME_BUDGET_SEC - elapsed()


def build_prompt(tokenizer, row) -> str:
    user_msg = USER_TEMPLATE.format(
        task_type=row["task_type"],
        eval_type=row["eval_type"],
        work_lang=row.get("work_lang", ""),
        task_lang=row.get("task_lang", ""),
        context=row["context"],
        query=row["query"],
    )
    messages = [
        {"role": "system", "content": SYSTEM},
        {"role": "user", "content": user_msg},
    ]
    return tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )


def adaptive_max_new_tokens(rows_left: int) -> int:
    if rows_left <= 0:
        return MIN_NEW_TOKENS
    per_row_budget = remaining_budget() / rows_left
    # Rough T4/AWQ-7B throughput assumption used only to size generation length;
    # errs conservative so we degrade gracefully rather than time out.
    tokens_per_sec_estimate = 20.0
    budget_tokens = int(per_row_budget * tokens_per_sec_estimate)
    return max(MIN_NEW_TOKENS, min(MAX_NEW_TOKENS, budget_tokens))


def adaptive_k(rows_left: int, max_new_tokens: int) -> int:
    """How many sampled completions per row the remaining budget affords.
    Only returns >1 once max_new_tokens has already saturated at
    MAX_NEW_TOKENS (i.e. there's more budget than a single generation pass
    needs) — never trades away reasoning length for extra samples.
    """
    if rows_left <= 0 or max_new_tokens <= 0:
        return MIN_K
    per_row_budget = remaining_budget() / rows_left
    tokens_per_sec_estimate = 20.0
    budget_tokens = int(per_row_budget * tokens_per_sec_estimate)
    k = budget_tokens // max_new_tokens
    return max(MIN_K, min(MAX_K, k))


def majority_vote_answers(sampled_answer_lists: list[list[str]]) -> list[str]:
    """Collapse k parsed answer lists (same length, one per sample) into one
    by majority vote per position. Prefers non-empty answers over blanks when
    both appear, since a wrong guess still earns partial chrF credit and a
    blank never does. Ties break toward the first sample's value.
    """
    if not sampled_answer_lists:
        return []
    n = len(sampled_answer_lists[0])
    result = []
    for i in range(n):
        votes = [sample[i] for sample in sampled_answer_lists if i < len(sample)]
        pool = [v for v in votes if v] or votes
        if not pool:
            result.append("")
            continue
        counts = Counter(pool)
        max_count = max(counts.values())
        winner = next(v for v in pool if counts[v] == max_count)
        result.append(winner)
    return result


def pad_row(expected_n: int) -> str:
    return json.dumps([""] * expected_n)


def main() -> None:
    df = pd.read_csv(DATA_PATH, dtype=str).fillna("")

    tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "left"

    model = AutoModelForCausalLM.from_pretrained(
        MODEL_PATH, torch_dtype=torch.float16, device_map="auto"
    )
    model.eval()

    results: dict[str, str] = {}
    rows = df.to_dict("records")

    i = 0
    while i < len(rows):
        rows_left = len(rows) - i

        if remaining_budget() < 30:  # not enough time left to safely attempt a batch
            for row in rows[i:]:
                expected_n = count_expected_items(row["query"])
                results[row["id"]] = pad_row(expected_n)
            break

        batch = rows[i : i + BATCH_SIZE]
        expected_counts = [count_expected_items(r["query"]) for r in batch]
        prompts = [build_prompt(tokenizer, r) for r in batch]

        max_new_tokens = adaptive_max_new_tokens(rows_left)
        k = adaptive_k(rows_left, max_new_tokens)

        inputs = tokenizer(
            prompts, return_tensors="pt", padding=True, truncation=True
        ).to(model.device)

        sampled_decoded: list[list[str]] = []
        for sample_idx in range(k):
            gen_kwargs = dict(
                max_new_tokens=max_new_tokens, pad_token_id=tokenizer.pad_token_id
            )
            if k > 1:
                gen_kwargs.update(do_sample=True, temperature=SAMPLE_TEMPERATURE, top_p=0.9)
            else:
                gen_kwargs.update(do_sample=False)
            try:
                with torch.no_grad():
                    output_ids = model.generate(**inputs, **gen_kwargs)
                input_len = inputs["input_ids"].shape[1]
                decoded = tokenizer.batch_decode(
                    output_ids[:, input_len:], skip_special_tokens=True
                )
            except Exception:
                decoded = [""] * len(batch)
            sampled_decoded.append(decoded)

        for row_idx, (row, expected_n) in enumerate(zip(batch, expected_counts)):
            sampled_answers = [
                parse_answers(sampled_decoded[s][row_idx], expected_n) for s in range(k)
            ]
            answers = majority_vote_answers(sampled_answers)
            results[row["id"]] = json.dumps(answers)

        i += len(batch)

    out_df = pd.DataFrame(
        {"id": df["id"], "pred": [results[i] for i in df["id"]]}
    )
    out_df.to_csv(OUTPUT_PATH, index=False)


if __name__ == "__main__":
    main()

	"""IOL-AI 2026 submission script.

	Runs fully offline against /tmp/data/test.csv and writes submission.csv with
	columns: id, pred (JSON list of answer strings), explanation (omitted).
	"""
	import os

	os.environ["HF_HUB_OFFLINE"] = "1"
	os.environ["TRANSFORMERS_OFFLINE"] = "1"

	import time

	T_START = time.monotonic()

	import json
	import re
	from collections import Counter

	import pandas as pd
	import torch
	from transformers import AutoModelForCausalLM, AutoTokenizer

	# --- Configuration -----------------------------------------------------

	# Production always loads from "." (weights shipped alongside this script in
	# the HF model repo). This constant only documents which Hub model those
	# weights came from, for keeping the Colab test notebook in sync.
	MODEL_ID_FOR_LOCAL_TEST = "Qwen/Qwen2.5-7B-Instruct-AWQ"
	# Fallback if Colab timing shows 7B has slack to spare: swap MODEL_ID_FOR_LOCAL_TEST
	# to "Qwen/Qwen2.5-14B-Instruct-AWQ" and re-run the notebook before re-uploading.

	MODEL_PATH = "."
	DATA_PATH = "/tmp/data/test.csv"
	OUTPUT_PATH = "submission.csv"

	TIME_BUDGET_SEC = 27 * 60 # 3-min safety margin under the 30-min hard limit
	BATCH_SIZE = 4
	MIN_NEW_TOKENS = 128
	MAX_NEW_TOKENS = 1024

	# Self-consistency: sample k >= 1 completions per row and majority-vote per
	# answer position when the time budget has slack. k=1 is plain greedy decode
	# (current default, deterministic). See adaptive_k().
	MIN_K = 1
	MAX_K = 1
	SAMPLE_TEMPERATURE = 0.7

	# Which SYSTEM prompt variant to use — "full_cot" or "plain_io". Published
	# evidence on whether heavy chain-of-thought helps or hurts a 7B-class model
	# on this task type is mixed (modeLing, arXiv:2406.17038), so both variants
	# are kept and should be A/B tested on the Colab harness rather than assumed.
	PROMPT_VARIANT = "plain_io"

	# Worked example shared by both prompt variants: a made-up toy language (not
	# a real one) demonstrating (a) the answer must be a word-form in the target
	# language, never an English gloss, and (b) subject/object roles must be
	# re-derived per query, not copied from a similar-looking example. Directly
	# targets the two failure modes observed in real Colab runs (language
	# confusion — arXiv:2406.20052 shows few-shot examples largely eliminate it;
	# and swapped argument roles).
	_WORKED_EXAMPLE = """Worked example (a made-up toy language, unrelated to your actual task — the
	words and rules below apply ONLY to this example; never reuse them):

	CONTEXT:
	mota = "dog" mota-ta = "dog" (as object)
	suno = "cat" suno-ta = "cat" (as object)
	kire = "sees"
	mota suno-ta kire = "the dog sees the cat"
	suno mota-ta kire = "the cat sees the dog"

	QUERY:
	Translate into the made-up language: "the cat sees the dog"

	IDENTIFY PATTERNS: word order is SUBJECT OBJECT-ta VERB; subject is unmarked,
	object takes suffix "-ta", verb is always last.
	RULE TABLE:
	- subject noun -> bare noun, placed first
	- object noun -> noun + "-ta", placed second
	- verb -> placed last, unchanged
	TEST: "mota suno-ta kire" = dog(subj) cat-ta(obj) sees = "the dog sees the
	cat" — matches. "suno mota-ta kire" = cat(subj) dog-ta(obj) sees = "the cat
	sees the dog" — matches.
	APPLY: query is "the cat sees the dog" -> subject=cat=suno,
	object=dog=mota-ta, verb=kire.
	VERIFY: answer is a word-form in the made-up language, not an English gloss.
	Roles checked: cat is subject (first, unmarked), dog is object (second,
	"-ta" suffix) — matches "the cat sees the dog", not reversed.

	FINAL ANSWERS:
	suno mota-ta kire

	--- end worked example. Now solve the actual task below using ONLY the data
	in its own CONTEXT and QUERY — never reuse the words or rules above. ---"""

	_ANSWER_FORMAT_RULES = """Answer formatting rules by task_type:
	- translation: give the full translated phrase or sentence for each numbered item,
	written as a natural fluent sentence exactly like the answer column in the given
	examples — never as a morpheme-by-morpheme gloss with parentheses like
	"you(shuddered)" or "we(spat(on him))".
	- fill_blanks: give only the missing word(s)/form for each numbered blank.
	- match_letters: give the matching letter or number for each item.
	- text_to_num: give only the numeral for each item.
	- num_to_text: give only the word(s) for each item.

	Output a line that says exactly:
	FINAL ANSWERS:
	followed by one answer per line, in the same order and count as the numbered
	items in the query, with no numbering, quotes, or extra commentary — just the
	bare answer text for each line."""

	SYSTEM_FULL_COT = f"""You are an expert linguist solving International Linguistics Olympiad problems.
	You are given a self-contained set of data in a language you have never seen before,
	plus a query asking you to apply what that data teaches you. You must reason only from
	the data given — never from memorized knowledge of real-world languages.

	Work in stages, showing your work:
	1. IDENTIFY PATTERNS: look for recurring forms, affixes, word order, sound
	correspondences, or structural regularities in the given data.
	2. RULE TABLE: write your hypotheses as short "TRIGGER -> TRANSFORMATION"
	bullet lines (e.g. "subject noun -> bare noun, placed first"), not prose —
	this keeps rules mechanical and easy to re-check, rather than vague
	descriptions that drift away from the actual data.
	3. TEST HYPOTHESES: check each rule-table line against every example in the
	given data. Discard or refine any rule that doesn't hold up on every
	example.
	4. APPLY: use only the rules that survived testing to answer the query.
	5. VERIFY: before writing FINAL ANSWERS, re-check every draft answer against
	two common mistakes:
	- LANGUAGE: look at the answer column in the given examples, not the
	instructions, to see what language/form your answer must be in. If the
	examples' answers are word-forms in the unfamiliar language, your answer
	must also be a word-form in that language — never substitute an English
	gloss or description of the meaning, even if you're unsure of the exact
	form; give your best-guess constructed form instead.
	- ROLES: for anything involving "who did what to whom" (subject, object,
	possessor, giver/receiver), re-read the query and re-derive which
	argument fills which role from scratch. Do not assume the same word
	order or role assignment as a similar-looking training example — verify
	it against the actual affixes/markers in that example.

	{_WORKED_EXAMPLE}

	{_ANSWER_FORMAT_RULES}"""

	SYSTEM_PLAIN_IO = f"""You are an expert linguist solving International Linguistics Olympiad problems.
	You are given a self-contained set of data in a language you have never seen before,
	plus a query asking you to apply what that data teaches you. Reason only from the
	data given — never from memorized knowledge of real-world languages. Give your
	best-guess answer in the same language/form as the examples' answers — never
	substitute an English gloss. Double-check which argument is subject vs object
	before answering.

	{_WORKED_EXAMPLE}

	{_ANSWER_FORMAT_RULES}"""

	_SYSTEM_VARIANTS = {"full_cot": SYSTEM_FULL_COT, "plain_io": SYSTEM_PLAIN_IO}
	SYSTEM = _SYSTEM_VARIANTS[PROMPT_VARIANT]

	USER_TEMPLATE = """Task type: {task_type}
	Eval type: {eval_type}
	Working language -> Task language: {work_lang} -> {task_lang}

	CONTEXT:
	{context}

	QUERY:
	{query}"""


	# --- Query parsing helpers ---------------------------------------------

	_RANGE_RE = re.compile(r"\((\d+)\s[-–]\s(\d+)\)")
	_LEADING_NUM_RE = re.compile(r"^\s*(\d+)[.)]", re.MULTILINE)
	_INLINE_PAREN_NUM_RE = re.compile(r"\((\d+)\)")


	def _is_contiguous(nums: list[int]) -> bool:
	if not nums:
	return False
	uniq = sorted(set(nums))
	return uniq == list(range(uniq[0], uniq[-1] + 1))


	def count_expected_items(query: str) -> int:
	"""Count how many answers the query expects, robust to numbering style."""
	range_match = _RANGE_RE.search(query)
	if range_match:
	start, end = int(range_match.group(1)), int(range_match.group(2))
	if end >= start:
	return end - start + 1

	leading_nums = [int(n) for n in _LEADING_NUM_RE.findall(query)]
	if _is_contiguous(leading_nums):
	return len(set(leading_nums))

	inline_nums = [int(n) for n in _INLINE_PAREN_NUM_RE.findall(query)]
	if _is_contiguous(inline_nums):
	return len(set(inline_nums))

	non_empty_lines = [ln for ln in query.splitlines() if ln.strip()]
	return max(len(non_empty_lines) - 1, 1)


	# --- Answer parsing helpers ----------------------------------------------

	_MARKER_RE = re.compile(r"FINAL ANSWERS:\s*", re.IGNORECASE)
	_LINE_NUM_PREFIX_RE = re.compile(r"^\s\(?\d+\)?[.):]?\s")
	_QUOTE_STRIP_RE = re.compile(r'^["\'](.*)["\']$')


	def _strip_line(line: str) -> str:
	line = line.strip()
	line = _LINE_NUM_PREFIX_RE.sub("", line, count=1)
	m = _QUOTE_STRIP_RE.match(line)
	if m:
	line = m.group(1)
	return line.strip()


	def parse_answers(text: str, expected_n: int) -> list[str]:
	"""Extract exactly expected_n answer strings from raw model output."""
	parts = _MARKER_RE.split(text, maxsplit=1)
	tail = parts[1] if len(parts) > 1 else text

	lines = [_strip_line(ln) for ln in tail.splitlines()]
	answers = [ln for ln in lines if ln]

	# Model sometimes emits all answers on one "\|"-delimited line instead of
	# one per line; only fall back to splitting on "\|" when the plain
	# line-split came up short, so genuine single-line answers containing "\|"
	# aren't mangled.
	if len(answers) < expected_n and any("\|" in a for a in answers):
	expanded = []
	for a in answers:
	if "\|" in a:
	expanded.extend(_strip_line(p) for p in a.split("\|"))
	else:
	expanded.append(a)
	expanded = [a for a in expanded if a]
	if len(expanded) > len(answers):
	answers = expanded

	if len(answers) < expected_n:
	answers = answers + [""] * (expected_n - len(answers))
	elif len(answers) > expected_n:
	answers = answers[:expected_n]

	return answers


	# --- Time-budget-aware batched generation --------------------------------

	def elapsed() -> float:
	return time.monotonic() - T_START


	def remaining_budget() -> float:
	return TIME_BUDGET_SEC - elapsed()


	def build_prompt(tokenizer, row) -> str:
	user_msg = USER_TEMPLATE.format(
	task_type=row["task_type"],
	eval_type=row["eval_type"],
	work_lang=row.get("work_lang", ""),
	task_lang=row.get("task_lang", ""),
	context=row["context"],
	query=row["query"],
	)
	messages = [
	{"role": "system", "content": SYSTEM},
	{"role": "user", "content": user_msg},
	]
	return tokenizer.apply_chat_template(
	messages, tokenize=False, add_generation_prompt=True
	)


	def adaptive_max_new_tokens(rows_left: int) -> int:
	if rows_left <= 0:
	return MIN_NEW_TOKENS
	per_row_budget = remaining_budget() / rows_left
	# Rough T4/AWQ-7B throughput assumption used only to size generation length;
	# errs conservative so we degrade gracefully rather than time out.
	tokens_per_sec_estimate = 20.0
	budget_tokens = int(per_row_budget * tokens_per_sec_estimate)
	return max(MIN_NEW_TOKENS, min(MAX_NEW_TOKENS, budget_tokens))


	def adaptive_k(rows_left: int, max_new_tokens: int) -> int:
	"""How many sampled completions per row the remaining budget affords.

	Only returns >1 once max_new_tokens has already saturated at
	MAX_NEW_TOKENS (i.e. there's more budget than a single generation pass
	needs) — never trades away reasoning length for extra samples.
	"""
	if rows_left <= 0 or max_new_tokens <= 0:
	return MIN_K
	per_row_budget = remaining_budget() / rows_left
	tokens_per_sec_estimate = 20.0
	budget_tokens = int(per_row_budget * tokens_per_sec_estimate)
	k = budget_tokens // max_new_tokens
	return max(MIN_K, min(MAX_K, k))


	def majority_vote_answers(sampled_answer_lists: list[list[str]]) -> list[str]:
	"""Collapse k parsed answer lists (same length, one per sample) into one
	by majority vote per position. Prefers non-empty answers over blanks when
	both appear, since a wrong guess still earns partial chrF credit and a
	blank never does. Ties break toward the first sample's value.
	"""
	if not sampled_answer_lists:
	return []
	n = len(sampled_answer_lists[0])
	result = []
	for i in range(n):
	votes = [sample[i] for sample in sampled_answer_lists if i < len(sample)]
	pool = [v for v in votes if v] or votes
	if not pool:
	result.append("")
	continue
	counts = Counter(pool)
	max_count = max(counts.values())
	winner = next(v for v in pool if counts[v] == max_count)
	result.append(winner)
	return result


	def pad_row(expected_n: int) -> str:
	return json.dumps([""] * expected_n)


	def main() -> None:
	df = pd.read_csv(DATA_PATH, dtype=str).fillna("")

	tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
	if tokenizer.pad_token is None:
	tokenizer.pad_token = tokenizer.eos_token
	tokenizer.padding_side = "left"

	model = AutoModelForCausalLM.from_pretrained(
	MODEL_PATH, torch_dtype=torch.float16, device_map="auto"
	)
	model.eval()

	results: dict[str, str] = {}
	rows = df.to_dict("records")

	i = 0
	while i < len(rows):
	rows_left = len(rows) - i

	if remaining_budget() < 30: # not enough time left to safely attempt a batch
	for row in rows[i:]:
	expected_n = count_expected_items(row["query"])
	results[row["id"]] = pad_row(expected_n)
	break

	batch = rows[i : i + BATCH_SIZE]
	expected_counts = [count_expected_items(r["query"]) for r in batch]
	prompts = [build_prompt(tokenizer, r) for r in batch]

	max_new_tokens = adaptive_max_new_tokens(rows_left)
	k = adaptive_k(rows_left, max_new_tokens)

	inputs = tokenizer(
	prompts, return_tensors="pt", padding=True, truncation=True
	).to(model.device)

	sampled_decoded: list[list[str]] = []
	for sample_idx in range(k):
	gen_kwargs = dict(
	max_new_tokens=max_new_tokens, pad_token_id=tokenizer.pad_token_id
	)
	if k > 1:
	gen_kwargs.update(do_sample=True, temperature=SAMPLE_TEMPERATURE, top_p=0.9)
	else:
	gen_kwargs.update(do_sample=False)
	try:
	with torch.no_grad():
	output_ids = model.generate(inputs, gen_kwargs)
	input_len = inputs["input_ids"].shape[1]
	decoded = tokenizer.batch_decode(
	output_ids[:, input_len:], skip_special_tokens=True
	)
	except Exception:
	decoded = [""] * len(batch)
	sampled_decoded.append(decoded)

	for row_idx, (row, expected_n) in enumerate(zip(batch, expected_counts)):
	sampled_answers = [
	parse_answers(sampled_decoded[s][row_idx], expected_n) for s in range(k)
	]
	answers = majority_vote_answers(sampled_answers)
	results[row["id"]] = json.dumps(answers)

	i += len(batch)

	out_df = pd.DataFrame(
	{"id": df["id"], "pred": [results[i] for i in df["id"]]}
	)
	out_df.to_csv(OUTPUT_PATH, index=False)


	if __name__ == "__main__":
	main()