Create script.py

script.py

@@ -0,0 +1,609 @@
+# script.py — Qwen3-14B-AWQ + OUR decomposition pipeline (Srikar's offline
+# wheelhouse used ONLY as the loading mechanism, not his prompt).
+# =============================================================================
+# WHAT CHANGED vs the 0.104 baseline (single conceptual variable = the model):
+#   1. Base model: Qwen2.5-14B-bnb-4bit  ->  Qwen3-14B-AWQ.
+#   2. Dependency install: instead of `pip install --no-deps bitsandbytes`
+#      (base env), we install the Qwen3 stack from BUNDLED wheels with
+#      `pip install --no-index --no-deps --target <RUNTIME_DIR>` and prepend
+#      that dir to sys.path. This NEVER touches the network and NEVER mutates
+#      base site-packages, so the grader's later hf_hub_download (metric.py)
+#      runs on the pristine base huggingface_hub -- the RemoteDisconnected
+#      class of failure cannot recur.
+#   3. Chat template: pass enable_thinking=False (Qwen3 supports it; harmless
+#      on models that ignore it). We keep OUR own decomposition reasoning in
+#      the prompt rather than paying for Qwen3's <think> phase.
+# Everything else -- symbolic evidence, FINAL ANSWERS contract, safe
+# arithmetic, explanations, per-row crash safety, dynamic token budget,
+# guaranteed one row per id -- is IDENTICAL to the proven 0.104 pipeline.
+# match_letters bijection decoding is deliberately NOT added here; that is the
+# next, separate experiment.
+# =============================================================================
+import os
+import atexit
+from pathlib import Path
+_ORIGINAL_HF_HUB_OFFLINE = os.environ.get("HF_HUB_OFFLINE")
+_ORIGINAL_TRANSFORMERS_OFFLINE = os.environ.get("TRANSFORMERS_OFFLINE")
+def _restore_offline_env_vars():
+    for key, original in (("HF_HUB_OFFLINE", _ORIGINAL_HF_HUB_OFFLINE),
+                           ("TRANSFORMERS_OFFLINE", _ORIGINAL_TRANSFORMERS_OFFLINE)):
+        if original is None:
+            os.environ.pop(key, None)
+        else:
+            os.environ[key] = original
+atexit.register(_restore_offline_env_vars)
+os.environ.setdefault("HF_HUB_OFFLINE", "1")
+os.environ.setdefault("TRANSFORMERS_OFFLINE", "1")
+import subprocess, sys
+import importlib
+import importlib.metadata
+SCRIPT_DIR = Path(__file__).resolve().parent
+def emergency_submission_csv(reason, rows_so_far=None):
+    try:
+        import pandas as pd
+        if rows_so_far:
+            pd.DataFrame(rows_so_far).to_csv("submission.csv", index=False)
+            return
+        try:
+            df = pd.read_csv("/tmp/data/test.csv", dtype=str).fillna("")
+            ids = df["id"].tolist()
+        except Exception:
+            ids = []
+        import json as _json
+        rows = [{"id": i, "pred": _json.dumps([""]),
+                  "explanation": f"EMERGENCY FALLBACK: {str(reason)[:150]}"} for i in ids]
+        pd.DataFrame(rows, columns=["id", "pred", "explanation"]).to_csv("submission.csv", index=False)
+    except Exception:
+        try:
+            with open("submission.csv", "w") as f:
+                f.write("id,pred,explanation\n")
+        except Exception:
+            pass
+def write_submission_csv(rows_list):
+    import csv as _csv
+    tmp_path = "submission.csv.tmp"
+    with open(tmp_path, "w", newline="", encoding="utf-8") as f:
+        w = _csv.DictWriter(f, fieldnames=["id", "pred", "explanation"])
+        w.writeheader()
+        for row in rows_list:
+            w.writerow(row)
+    os.replace(tmp_path, "submission.csv")
+# =============================================================================
+# OFFLINE WHEELHOUSE INSTALL (adapted from the public 0.147 submissions).
+# Installs the Qwen3-compatible stack from wheels bundled inside this repo,
+# to an isolated --target dir that we prepend to sys.path. --no-index means
+# pip never contacts the network (the sandbox has no working index anyway);
+# --target means base site-packages is untouched, so scoring stays safe.
+# =============================================================================
+WHEELHOUSE = Path(os.environ.get("QWEN3_WHEELHOUSE", str(SCRIPT_DIR / "wheelhouse")))
+RUNTIME_DIR = Path(os.environ.get("QWEN3_RUNTIME_DIR", "/tmp/qwen3deps"))
+RUNTIME_PACKAGES = {
+    "transformers": "4.51.3",
+    "tokenizers": "0.21.1",
+    "huggingface_hub": "0.30.2",
+    "autoawq": "0.2.9",
+}
+RUNTIME_WHEELS = (
+    "transformers-4.51.3-py3-none-any.whl",
+    "tokenizers-0.21.1-cp39-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl",
+    "huggingface_hub-0.30.2-py3-none-any.whl",
+    "autoawq-0.2.9-py3-none-any.whl",
+)
+def ensure_runtime_dependencies():
+    wheel_paths = [WHEELHOUSE / name for name in RUNTIME_WHEELS]
+    missing = [str(p) for p in wheel_paths if not p.is_file()]
+    if missing:
+        raise FileNotFoundError(f"Missing offline runtime wheels: {missing}")
+    marker = RUNTIME_DIR / ".iol-qwen3-runtime-v1"
+    if not marker.is_file():
+        RUNTIME_DIR.mkdir(parents=True, exist_ok=True)
+        subprocess.run(
+            [sys.executable, "-m", "pip", "install",
+             "--disable-pip-version-check", "--no-index", "--no-deps",
+             "--upgrade", "--target", str(RUNTIME_DIR),
+             *(str(p) for p in wheel_paths)],
+            check=True, timeout=300,
+        )
+        marker.write_text("offline Qwen3 runtime installed\n", encoding="utf-8")
+    runtime_path = str(RUNTIME_DIR)
+    if runtime_path in sys.path:
+        sys.path.remove(runtime_path)
+    sys.path.insert(0, runtime_path)
+    importlib.invalidate_caches()
+    versions = {}
+    for pkg in RUNTIME_PACKAGES:
+        try:
+            versions[pkg] = importlib.metadata.version(pkg)
+        except importlib.metadata.PackageNotFoundError:
+            versions[pkg] = "missing"
+    print(f"offline runtime active: {versions}", flush=True)
+try:
+    ensure_runtime_dependencies()
+except Exception as e:
+    emergency_submission_csv(f"wheelhouse install failed: {e}")
+    raise
+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 = 360
+start_time = time.time()
+try:
+    df = pd.read_csv("/tmp/data/test.csv", dtype=str).fillna("")
+    placeholder_rows = [{"id": rid, "pred": json.dumps([""]),
+                          "explanation": "Placeholder written before model load."}
+                         for rid in df["id"].tolist()]
+    write_submission_csv(placeholder_rows)
+    print(f"Pre-load checkpoint written for {len(placeholder_rows)} rows.", flush=True)
+    # AWQ backend preflight (diagnostic only, never fatal).
+    try:
+        from awq.modules.linear import gemm as awq_gemm
+        print(f"AWQ backends: extension={awq_gemm.awq_ext is not None}, "
+              f"triton={getattr(awq_gemm, 'TRITON_AVAILABLE', None)}", flush=True)
+    except Exception as exc:
+        print(f"AWQ backend preflight warning: {exc}", flush=True)
+    try:
+        tok = AutoTokenizer.from_pretrained(MODEL_ID, local_files_only=True)
+        print("Tokenizer loaded (fast).", flush=True)
+    except Exception as e:
+        print(f"Fast tokenizer failed ({e}); falling back to use_fast=False.", flush=True)
+        tok = AutoTokenizer.from_pretrained(MODEL_ID, use_fast=False, local_files_only=True)
+        print("Tokenizer loaded (slow fallback).", flush=True)
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_ID, torch_dtype=torch.float16, device_map="auto", local_files_only=True,
+    ).eval()
+    print(f"Model loaded | memory footprint: {round(model.get_memory_footprint()/1e9, 1)} GB | "
+          f"quantized: {getattr(model.config, 'quantization_config', None) is not None}", flush=True)
+except Exception as e:
+    emergency_submission_csv(f"tokenizer/model load or test.csv read failed: {e}")
+    raise
+n_rows = len(df)
+actual_setup_elapsed = time.time() - start_time
+per_row_budget = max(20, (TIME_LIMIT_S - actual_setup_elapsed) / max(n_rows, 1))
+print(f"Setup took {actual_setup_elapsed:.0f}s | per_row_budget={per_row_budget:.0f}s "
+      f"for {n_rows} rows", flush=True)
+# ---- Query parsing ----
+def parse_items(query: str):
+    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*(?:[-–—:]|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."
+from difflib import SequenceMatcher
+from collections import defaultdict
+def extract_forms_from_context(context: str):
+    forms = []
+    for line in context.splitlines():
+        line = line.strip()
+        if not line:
+            continue
+        pipe_count = line.count("|")
+        if 0 < pipe_count <= 3:
+            first_field = re.sub(r"^\s*\d+\s*[.\)]\s*", "", line.split("|")[0].strip()).strip()
+            if first_field:
+                forms.append(first_field)
+        elif pipe_count == 0:
+            for t in line.split():
+                t_clean = re.sub(r"^\s*\d+\s*[.\)]\s*", "", t).strip(".,;:")
+                if t_clean and len(t_clean) > 1:
+                    forms.append(t_clean)
+    seen, unique_forms = set(), []
+    for f in forms:
+        if f not in seen:
+            seen.add(f)
+            unique_forms.append(f)
+    return unique_forms
+def extract_explicit_pairs(context: str):
+    pairs = []
+    for line in context.splitlines():
+        line = line.strip()
+        if not (0 < line.count("|") <= 3):
+            continue
+        fields = [re.sub(r"^\s*\d+\s*[.\)]\s*", "", f.strip()).strip() for f in line.split("|")]
+        fields = [f for f in fields if f]
+        if len(fields) >= 2:
+            pairs.append((fields[0], fields[1]))
+    return pairs
+def edit_signature(a: str, b: str):
+    sm = SequenceMatcher(None, a, b, autojunk=False)
+    all_ops = sm.get_opcodes()
+    ops = [op for op in all_ops if op[0] != "equal"]
+    if not ops or len(ops) > 2:
+        return None
+    equal_len = sum((i2 - i1) for tag, i1, i2, j1, j2 in all_ops if tag == "equal")
+    if equal_len < 2:
+        return None
+    tag, i1, i2, j1, j2 = ops[0]
+    removed, inserted = a[i1:i2], b[j1:j2]
+    if i1 == 0:
+        pos = "prefix"
+    elif i2 == len(a):
+        pos = "suffix"
+    else:
+        pos = "infix"
+    return (pos, removed, inserted)
+def find_transformation_families(pairs):
+    groups = defaultdict(list)
+    for a, b in pairs:
+        if not a or not b or a == b:
+            continue
+        sig = edit_signature(a, b)
+        if sig:
+            groups[sig].append((a, b))
+    families = []
+    for sig, grp in groups.items():
+        unique_pairs = list(dict.fromkeys(grp))
+        if len(unique_pairs) >= 2:
+            pos, removed, inserted = sig
+            removed_disp = removed if removed else "(nothing)"
+            inserted_disp = inserted if inserted else "(nothing)"
+            examples = "; ".join(f"{a}->{b}" for a, b in unique_pairs[:4])
+            families.append((len(unique_pairs),
+                              f"{pos} change: '{removed_disp}' -> '{inserted_disp}' (seen in: {examples})"))
+    families.sort(key=lambda x: -x[0])
+    return [f for _, f in families]
+def detect_reduplication(forms):
+    findings = []
+    for w in forms:
+        n = len(w)
+        found = False
+        for length in range(2, n // 2 + 1):
+            for start in range(0, n - 2 * length + 1):
+                chunk = w[start:start + length]
+                nxt = w[start + length:start + 2 * length]
+                if chunk == nxt:
+                    findings.append(f"reduplication in '{w}': '{chunk}' repeated")
+                    found = True
+                    break
+            if found:
+                break
+    return findings
+def build_symbolic_evidence(context: str) -> str:
+    forms = extract_forms_from_context(context)
+    pairs = extract_explicit_pairs(context)
+    families = find_transformation_families(pairs) if pairs else []
+    redup = detect_reduplication(forms) if forms else []
+    lines = []
+    if families:
+        lines.append("Transformation families found (patterns supported by multiple examples):")
+        for f in families[:3]:
+            lines.append(f"- {f}")
+    if redup:
+        lines.append("Reduplication detected:")
+        for r in redup[:2]:
+            lines.append(f"- {r}")
+    if not lines:
+        return ""
+    return ("\n\nSYMBOLIC EVIDENCE (deterministically computed from the examples above; "
+            "may be incomplete -- verify against the examples, do not trust blindly):\n"
+            + "\n".join(lines))
+def build_messages(context, query, task_type):
+    preamble, items, count_known = parse_items(query)
+    guidance = TASK_GUIDANCE.get(task_type, DEFAULT_GUIDANCE)
+    symbolic_evidence = build_symbolic_evidence(context)
+    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."
+        )
+    options_note = ""
+    if task_type == "match_letters":
+        options = extract_match_letter_options(context)
+        if options:
+            options_note = (
+                f"\n\nThe only valid answers are: {', '.join(options)}. "
+                f"Do not use any other letter."
+            )
+    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()}"
+            f"{symbolic_evidence}\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}"
+            f"{options_note}"
+        )
+    else:
+        n_items = None
+        user = (
+            f"EXAMPLES:\n{context.strip()}"
+            f"{symbolic_evidence}\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}"
+            f"{options_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}]
+_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*(the\s+)?(final\s+)?answer\s*\d*\s*(is)?\s*:\s*", "", a).strip()
+    a = re.sub(r"(?i)^\s*is\s*:\s*", "", a).strip()
+    a = a.strip("* ")
+    return a.strip(" .\"'“”‘’")
+def extract(text):
+    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()]
+    fallback = []
+    for ln in lines:
+        ln_clean = re.sub(r"^\s*\d+\s*[.\)]\s*", "", ln)
+        if "|" in ln_clean:
+            fallback.extend(clean_answer(p) for p in ln_clean.split("|") if p.strip())
+        else:
+            fallback.append(clean_answer(ln_clean))
+    return fallback, 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. enable_thinking=False keeps Qwen3 in its fast, non-<think>
+# mode; our decomposition prompt supplies the reasoning instead. The kwarg is
+# harmless on templates that ignore it. Both API-shape branches pass it. ----
+def generate(messages, max_new_tokens, constraint_fn=None):
+    def _try_generate(gen_kwargs):
+        try:
+            enc = tok.apply_chat_template(
+                messages, add_generation_prompt=True, enable_thinking=False,
+                return_tensors="pt", return_dict=True,
+            ).to(model.device)
+            input_len = enc["input_ids"].shape[-1]
+            with torch.no_grad():
+                out = model.generate(**enc, **gen_kwargs)
+        except Exception:
+            ids = tok.apply_chat_template(
+                messages, add_generation_prompt=True, enable_thinking=False,
+                return_tensors="pt",
+            ).to(model.device)
+            input_len = ids.shape[-1]
+            with torch.no_grad():
+                out = model.generate(ids, **gen_kwargs)
+        return out, input_len
+    base_kwargs = {"max_new_tokens": max_new_tokens, "do_sample": False}
+    if constraint_fn is not None:
+        try:
+            out, input_len = _try_generate({**base_kwargs, "prefix_allowed_tokens_fn": constraint_fn})
+        except Exception:
+            out, input_len = _try_generate(base_kwargs)
+    else:
+        out, input_len = _try_generate(base_kwargs)
+    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."
+_LETTER_CONSTRAINT_CACHE = {}
+def build_letter_constraint_fn(tok, valid_letters):
+    cache_key = (id(tok), tuple(sorted(valid_letters)))
+    if cache_key in _LETTER_CONSTRAINT_CACHE:
+        return _LETTER_CONSTRAINT_CACHE[cache_key]
+    try:
+        allowed_chars = set(valid_letters) | set(" |\n\t\r")
+        eos = tok.eos_token_id
+        pieces = []
+        for token_id in range(len(tok)):
+            if token_id == eos:
+                continue
+            piece = tok.decode([token_id], skip_special_tokens=False)
+            if piece and all(c in allowed_chars for c in piece):
+                pieces.append(token_id)
+        allowed_ids = ([eos] if eos is not None else []) + pieces
+        def allowed(_batch_id, _input_ids):
+            return allowed_ids if allowed_ids else list(range(len(tok)))
+        _LETTER_CONSTRAINT_CACHE[cache_key] = allowed
+        return allowed
+    except Exception:
+        return None
+def extract_match_letter_options(context: str):
+    found = set()
+    for line in context.splitlines():
+        for m in re.finditer(r"(?:^|\s)([A-Z])[.\)]\s+\S", line):
+            found.add(m.group(1))
+    if not found:
+        return None
+    letters = sorted(found)
+    expected = [chr(ord("A") + i) for i in range(len(letters))]
+    if letters != expected:
+        return None
+    if not (2 <= len(letters) <= 26):
+        return None
+    return letters
+rows = []
+processed_ids = set()
+try:
+    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
+            time_based_cap = 1280 if row_budget > per_row_budget else 640
+            task_type = r.get("task_type", "")
+            messages, n_items = build_messages(r["context"], r["query"], task_type)
+            if n_items:
+                item_based_cap = max(640, min(1536, n_items * 48 + 256))
+                tokens_cap = min(time_based_cap, item_based_cap)
+            else:
+                tokens_cap = time_based_cap
+            text = generate(messages, tokens_cap)
+            answers, marker_pos = extract(text)
+            if 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
+            if (marker_pos is None or not answers) and remaining > SETUP_BUFFER_S:
+                repair_constraint = None
+                if task_type == "match_letters":
+                    repair_options = extract_match_letter_options(r["context"])
+                    if repair_options:
+                        repair_constraint = build_letter_constraint_fn(tok, repair_options)
+                repair_text = generate(build_repair_messages(r["query"], n_items, text), 128,
+                                        constraint_fn=repair_constraint)
+                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]
+            if not answers:
+                answers = [""]
+            remaining_after = TIME_LIMIT_S - (time.time() - start_time)
+            budget_left_after = max(n_rows - len(rows) - 1, 0)
+            comfortable = remaining_after > (budget_left_after + 1) * per_row_budget * 1.3
+            if comfortable:
+                try:
+                    explanation = generate(
+                        [{"role": "system", "content": EXPLANATION_SYSTEM},
+                         {"role": "user", "content": text}], 300,
+                    ) or EXPLANATION_FALLBACK
+                except Exception:
+                    explanation = EXPLANATION_FALLBACK
+            else:
+                snippet = re.sub(r"\s{2,}", " ", text[:300]).strip()
+                explanation = snippet if snippet else EXPLANATION_FALLBACK
+            rows.append({"id": r["id"], "pred": json.dumps(answers, ensure_ascii=False),
+                         "explanation": explanation})
+            processed_ids.add(r["id"])
+            write_submission_csv(rows)
+            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"])
+            write_submission_csv(rows)
+            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
+    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})
+    write_submission_csv(rows)
+    print("DONE.", flush=True)
+except Exception as e:
+    emergency_submission_csv(f"main loop failed: {e}", rows_so_far=rows if rows else None)
+    print(f"FATAL, but submission.csv was written with {len(rows)} rows. Error: {e}", flush=True)