script.py

# 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)