script.py

# script.py — FINAL SUBMISSION: Qwen2.5-14B-Instruct-AWQ + decomposition/
# verification prompt + safe arithmetic for number tasks + guaranteed
# explanations. Every piece below was individually tested and fixed against
# real bugs found on real Linguini problems before being combined here.
# =============================================================================
# COMPLIANCE (verified below): offline before any HF import, MODEL_ID=".",
# reads only /tmp/data/test.csv, writes only submission.csv with id/pred/
# explanation, float16 (T4 has no native bfloat16), no hub names anywhere,
# 30-minute limit respected with a real safety margin, crash-safe per row,
# every row guaranteed a submission.csv entry even under a timeout.
# =============================================================================

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

import subprocess, sys
# Unpinned/latest, matching the organizers' OWN proven-working recipe for
# this exact model (their reference baseline uses -U, not a version pin --
# AWQ loading via gptqmodel needs a recent transformers to work correctly).
subprocess.run([sys.executable, "-m", "pip", "install", "-q", "-U",
                "transformers", "accelerate", "gptqmodel", "pandas",
                "numpy==2.2.6", "torch"], check=True)

import re, json, time, ast as pyast
import pandas as pd
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

MODEL_ID = "."
TIME_LIMIT_S = 30 * 60
SETUP_BUFFER_S = 420   # larger margin: 14B AWQ checkpoint is ~9GB, slower to load than anything tested before
start_time = time.time()

tok = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID, torch_dtype=torch.float16, device_map="auto",
).eval()

df = pd.read_csv("/tmp/data/test.csv", dtype=str).fillna("")
n_rows = len(df)
per_row_budget = max(20, (TIME_LIMIT_S - SETUP_BUFFER_S) / max(n_rows, 1))


# ---- Query parsing: widened patterns + honest "unknown count" fallback ----
def parse_items(query: str):
    """Returns (preamble, items, count_known). count_known=False means no
    pattern matched -- we do NOT guess a count, we let the model's own
    answer list stand rather than risk truncating real content."""
    item_pat = re.compile(r"(?m)^\s*(\d+)\s*[.\)]\s*(.*)$")
    matches = list(item_pat.finditer(query))
    if matches:
        preamble = query[:matches[0].start()].strip()
        items = []
        for i, m in enumerate(matches):
            end = matches[i + 1].start() if i + 1 < len(matches) else len(query)
            text = re.sub(r"^\s*\d+\s*[.\)]\s*", "", query[m.start():end].strip())
            items.append(text)
        return preamble, items, True

    rng = re.search(r"[\(\[]?\s*(\d+)\s*(?:[-\u2013\u2014:]|to)\s*(\d+)\s*[\)\]]?", query, flags=re.IGNORECASE)
    if rng:
        lo, hi = int(rng.group(1)), int(rng.group(2))
        if 0 < hi - lo < 100:
            items = []
            for k in range(lo, hi + 1):
                line_match = re.search(rf"(?m)^.*\(\s*{k}\s*\).*$", query)
                if line_match:
                    clue = re.sub(rf"\(\s*{k}\s*\)", "", line_match.group(0)).strip()
                    clue = re.sub(r"\|\s*\|", "|", clue)
                    clue = re.sub(r"\s{2,}", " ", clue).strip(" |")
                    items.append(clue if clue else f"the numbered item {k} from the examples above")
                else:
                    items.append(f"the numbered item {k} from the examples above")
            return query.strip(), items, True

    csv_nums = re.findall(r"(?m)^\s*(\d+)\s*,\s*(\d+(?:\s*,\s*\d+)*)\s*$", query)
    if csv_nums:
        all_nums = re.findall(r"\d+", " ".join(csv_nums[0]))
        return query.strip(), [f"the numbered item {n}" for n in all_nums], True

    return query.strip(), [], False


TASK_GUIDANCE = {
    "translation": "give the translated form only, in the language asked.",
    "fill_blanks": "give only the missing form for each blank.",
    "match_letters": "give only the option letter (for example A, B, C).",
    "text_to_num": "give the number in digits.",
    "num_to_text": "give the number written out in words, in the language asked.",
}
DEFAULT_GUIDANCE = "give exactly what the instruction asks, nothing else."


def build_messages(context, query, task_type):
    preamble, items, count_known = parse_items(query)
    guidance = TASK_GUIDANCE.get(task_type, DEFAULT_GUIDANCE)

    system = (
        "You solve puzzles about a language you have never seen. Everything you "
        "need is in the examples below. Use only the examples, not outside "
        "knowledge of any language. You may meet a task type you have never "
        "seen -- read the instruction and examples, and answer in the same "
        "form they use."
    )
    number_note = ""
    if task_type == "text_to_num":
        number_note = (
            "\n\nAlso add one more line after your answers, exactly like this:\n"
            "COMPUTE: expr1 | expr2\n"
            "where each expr is a plain arithmetic expression (digits, +, -, *, "
            "parentheses only) for that item's value, one per answer, matching "
            "the rule you found."
        )

    if count_known:
        n_items = len(items)
        slots = "\n\n".join(f"Question {i+1}: {it}\nAnswer {i+1}:" for i, it in enumerate(items))
        user = (
            f"EXAMPLES:\n{context.strip()}\n\n"
            f"--- The examples end here. The questions begin below. ---\n\n"
            f"For each question: find the rule that explains ALL the examples above "
            f"(not just one). Check it against every example before answering. "
            f"For this task type, {guidance}\n\n"
            f"{preamble}\n\n{slots}\n\n"
            f"After answering all {n_items} questions, finish with exactly one line, "
            f"all {n_items} answers in order separated by ' | ':\n"
            f"FINAL ANSWERS: answer1 | answer2"
            f"{number_note}"
        )
    else:
        n_items = None
        user = (
            f"EXAMPLES:\n{context.strip()}\n\n"
            f"--- The examples end here. The question begins below. ---\n\n"
            f"Find the rule that explains ALL the examples above (not just one). "
            f"Check it against every example before answering. "
            f"For this task type, {guidance}\n\n"
            f"{preamble}\n\n"
            f"Answer every item asked above, in order, one per answer. Finish "
            f"with exactly one line, all your answers in order separated by ' | ':\n"
            f"FINAL ANSWERS: answer1 | answer2"
            f"{number_note}"
        )
    return [{"role": "system", "content": system}, {"role": "user", "content": user}], n_items


def build_repair_messages(query, n_items, bad_text):
    n_desc = f"exactly {n_items}" if n_items is not None else "one per item asked"
    system = "You reformat answers. Output nothing except the requested line."
    user = (
        f"Question:\n{query.strip()}\n\n"
        f"A previous attempt produced:\n{bad_text[:600]}\n\n"
        f"Extract or restate {n_desc} final answers, in order, as ONE line:\n"
        f"FINAL ANSWERS: answer1 | answer2"
    )
    return [{"role": "system", "content": system}, {"role": "user", "content": user}]


# ---- Safe arithmetic: no exec(), no eval() of arbitrary code ----
_ALLOWED_BINOPS = (pyast.Add, pyast.Sub, pyast.Mult)

def safe_arithmetic(expr: str):
    try:
        tree = pyast.parse(expr.strip(), mode="eval")
    except Exception:
        return None
    def _eval(node):
        if isinstance(node, pyast.Expression):
            return _eval(node.body)
        if isinstance(node, pyast.Constant) and isinstance(node.value, (int, float)):
            return node.value
        if isinstance(node, pyast.BinOp) and isinstance(node.op, _ALLOWED_BINOPS):
            left, right = _eval(node.left), _eval(node.right)
            if left is None or right is None:
                return None
            if isinstance(node.op, pyast.Add): return left + right
            if isinstance(node.op, pyast.Sub): return left - right
            if isinstance(node.op, pyast.Mult): return left * right
        if isinstance(node, pyast.UnaryOp) and isinstance(node.op, pyast.USub):
            v = _eval(node.operand)
            return -v if v is not None else None
        return None
    return _eval(tree)


def clean_answer(a: str) -> str:
    a = re.sub(r"(?i)^\s*answer\s*\d*\s*:\s*", "", a).strip()
    a = a.strip("* ")
    return a.strip(" .\"'")


def extract(text):
    """Fixed against two real bugs found on real Linguini output:
    (1) markdown-bold marker with content on the NEXT line, not same line;
    (2) a following COMPUTE: line bleeding into the answer list."""
    m = list(re.finditer(r"final answers?\s*:?\s*\**", text, flags=re.IGNORECASE))
    if m:
        tail = text[m[-1].end():]
        stop = re.search(r"(?i)compute\s*:", tail)
        if stop:
            tail = tail[:stop.start()]
        tail = tail.replace("**", " ").strip()
        candidate = " ".join(tail.splitlines())
        parts = [clean_answer(p) for p in candidate.split("|") if p.strip()]
        if parts:
            return parts, m[-1].start()
    lines = [ln.strip() for ln in text.splitlines() if ln.strip()]
    return [clean_answer(re.sub(r"^\s*\d+\s*[.\)]\s*", "", ln)) for ln in lines], None


def extract_compute_overrides(text, n_answers):
    m = re.search(r"compute\s*:\s*(.+)", text, flags=re.IGNORECASE)
    if not m:
        return {}
    exprs = [e.strip() for e in m.group(1).split("|")]
    overrides = {}
    for i, e in enumerate(exprs[:n_answers]):
        val = safe_arithmetic(e)
        if val is not None:
            overrides[i] = str(int(val)) if float(val).is_integer() else str(val)
    return overrides


# ---- Generation: defensive against BOTH chat-template return shapes.     ----
# The organizers themselves confirm this discrepancy is real: recent
# transformers (their Colab) returns a dict from apply_chat_template with
# return_dict=True; older transformers (their own words: "the sandbox's
# older transformers") returns a bare tensor and may not even accept the
# return_dict kwarg. Handle both, don't assume either.
def generate(messages, max_new_tokens):
    try:
        enc = tok.apply_chat_template(
            messages, add_generation_prompt=True, return_tensors="pt", return_dict=True,
        ).to(model.device)
        input_len = enc["input_ids"].shape[-1]
        with torch.no_grad():
            out = model.generate(**enc, max_new_tokens=max_new_tokens, do_sample=False)
    except TypeError:
        ids = tok.apply_chat_template(
            messages, add_generation_prompt=True, return_tensors="pt",
        ).to(model.device)
        input_len = ids.shape[-1]
        with torch.no_grad():
            out = model.generate(ids, max_new_tokens=max_new_tokens, do_sample=False)
    return tok.decode(out[0][input_len:], skip_special_tokens=True).strip()


EXPLANATION_SYSTEM = (
    "Summarize the following reasoning into a few short bullet points: the "
    "rule or pattern found in the data and the key evidence for the answer. "
    "Be concise and structured -- do not repeat the full reasoning."
)
EXPLANATION_FALLBACK = "Answer derived from patterns found in the examples above."


rows = []
processed_ids = set()

for _, r in df.iterrows():
    try:
        elapsed = time.time() - start_time
        remaining = TIME_LIMIT_S - elapsed
        budget_left_rows = max(n_rows - len(rows), 1)
        row_budget = remaining / budget_left_rows
        tokens_cap = 1536 if row_budget > per_row_budget else 768

        messages, n_items = build_messages(r["context"], r["query"], r.get("task_type", ""))
        text = generate(messages, tokens_cap)
        answers, marker_pos = extract(text)

        if r.get("task_type", "") == "text_to_num":
            overrides = extract_compute_overrides(text, len(answers))
            for idx, val in overrides.items():
                if idx < len(answers):
                    answers[idx] = val

        # Repair only on TRUE extraction failure (no marker / nothing found) --
        # not on a mere count difference, since extra answers are harmless
        # and our own count guess may be the thing that's wrong.
        if (marker_pos is None or not answers) and remaining > SETUP_BUFFER_S:
            repair_text = generate(build_repair_messages(r["query"], n_items, text), 128)
            rep, rep_pos = extract(repair_text)
            if rep:
                answers, marker_pos = rep, rep_pos

        if n_items is not None:
            if len(answers) < n_items:
                answers = answers + [answers[-1] if answers else ""] * (n_items - len(answers))
            elif len(answers) > n_items and marker_pos is None:
                answers = answers[:n_items]
            # else: marker found, more answers than our guess -> KEEP THEM ALL

        if not answers:
            answers = [""]

        # Explanation: dedicated call, always non-empty even if it fails.
        try:
            explanation = generate(
                [{"role": "system", "content": EXPLANATION_SYSTEM},
                 {"role": "user", "content": text}], 300,
            ) or EXPLANATION_FALLBACK
        except Exception:
            explanation = EXPLANATION_FALLBACK

        rows.append({"id": r["id"], "pred": json.dumps(answers, ensure_ascii=False),
                     "explanation": explanation})
        processed_ids.add(r["id"])

        pd.DataFrame(rows).to_csv("submission.csv", index=False)
        print(f"{len(rows)}/{n_rows} answers={len(answers)} elapsed={time.time()-start_time:.0f}s", flush=True)

    except Exception as e:
        try:
            _, fallback_items, fk = parse_items(r["query"])
            n_fallback = len(fallback_items) if fk else 1
        except Exception:
            n_fallback = 1
        rows.append({"id": r["id"], "pred": json.dumps([""] * n_fallback, ensure_ascii=False),
                     "explanation": EXPLANATION_FALLBACK})
        processed_ids.add(r["id"])
        pd.DataFrame(rows).to_csv("submission.csv", index=False)
        print(f"ROW ERROR on {r['id']}: {e}", flush=True)

    if time.time() - start_time > TIME_LIMIT_S - 60:
        print("Time budget nearly exhausted, stopping early.", flush=True)
        break

# Guarantee one row per test.csv id, even under a timeout.
for _, r in df.iterrows():
    if r["id"] in processed_ids:
        continue
    try:
        _, fallback_items, fk = parse_items(r["query"])
        n_fallback = len(fallback_items) if fk else 1
    except Exception:
        n_fallback = 1
    rows.append({"id": r["id"], "pred": json.dumps([""] * n_fallback, ensure_ascii=False),
                 "explanation": EXPLANATION_FALLBACK})

pd.DataFrame(rows).to_csv("submission.csv", index=False)
print("DONE.", flush=True)