File size: 28,670 Bytes
4acf6ea
 
 
 
 
ef1cd02
 
 
 
 
 
 
 
 
 
 
 
 
6eea194
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
ef1cd02
 
 
 
 
 
 
 
4acf6ea
ef1cd02
 
6eea194
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
ef1cd02
 
6eea194
 
 
 
 
 
 
ef1cd02
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
afd21e9
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
ef1cd02
 
afd21e9
ef1cd02
 
afd21e9
 
 
 
 
ef1cd02
 
 
 
 
 
 
 
afd21e9
ef1cd02
 
 
 
 
 
afd21e9
 
 
 
 
ef1cd02
 
 
 
 
 
 
 
 
 
afd21e9
 
 
 
 
ef1cd02
 
 
 
 
 
 
 
 
 
d3bb616
 
ef1cd02
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
6eea194
 
 
 
 
ef1cd02
6eea194
ef1cd02
 
 
6eea194
ef1cd02
6eea194
 
 
 
ef1cd02
 
 
 
 
 
 
 
 
 
 
6eea194
 
ef1cd02
6eea194
 
 
 
 
 
 
 
ef1cd02
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
6eea194
 
 
 
 
 
ef1cd02
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
6eea194
 
 
 
 
 
 
afd21e9
d3bb616
 
afd21e9
 
6eea194
 
d3bb616
6eea194
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
afd21e9
 
 
6eea194
 
 
 
 
 
 
afd21e9
6eea194
 
 
 
afd21e9
6eea194
 
 
 
 
 
 
 
 
 
4acf6ea
6eea194
 
4acf6ea
6eea194
 
 
 
 
 
 
 
 
 
 
 
 
 
 
ef1cd02
 
 
 
 
 
 
 
6eea194
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
# script.py — FINAL SUBMISSION: Qwen2.5-14B-Instruct (bnb-4bit, via
# unsloth/Qwen2.5-14B-Instruct-bnb-4bit) + 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

def emergency_submission_csv(reason, rows_so_far=None):
    """Last-resort guarantee: no matter WHERE the script dies, write a valid
    submission.csv before the process exits. This is the single fix for the
    pattern behind both real failures so far -- a crash with nothing written
    produces the secondary 'not a file on the local file system' error every
    time, turning a scoreable zero into a hard evaluation failure."""
    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:
        # Absolute last resort: a header-only file is still a file.
        try:
            with open("submission.csv", "w") as f:
                f.write("id,pred,explanation\n")
        except Exception:
            pass


try:
    # Split deliberately: torch is NOT force-upgraded. It's a multi-GB
    # CUDA-specific wheel; forcing -U risks pulling a build mismatched with
    # the sandbox's actual driver -- a worse failure mode (silent GPU
    # incompatibility) than a missing package. bitsandbytes already
    # succeeded as-is in the last real run, no evidence it needs upgrading.
    # Only transformers/accelerate/tokenizers have a CONFIRMED version-
    # related failure behind them -- those are the only ones forced.
    subprocess.run([sys.executable, "-m", "pip", "install", "-q",
                    "torch>=2.2", "bitsandbytes", "pandas"], check=True)
    subprocess.run([sys.executable, "-m", "pip", "install", "-q", "-U",
                    "transformers>=4.43", "accelerate>=0.30", "tokenizers"], check=True)
except Exception as e:
    emergency_submission_csv(f"pip 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 = 420   # larger margin: 14B bnb-4bit checkpoint is ~8-9GB, slower to load than anything tested before
start_time = time.time()

try:
    try:
        tok = AutoTokenizer.from_pretrained(MODEL_ID)
        print("Tokenizer loaded (fast).", flush=True)
    except Exception as e:
        # Mechanism-level fix: bypasses TokenizerFast.from_file() entirely,
        # which is exactly the call that fails on a tokenizer.json saved by a
        # newer tokenizers library than the sandbox has. Falls back to the
        # pure Python tokenizer built from vocab.json/merges.txt instead.
        print(f"Fast tokenizer failed ({e}); falling back to use_fast=False.", flush=True)
        tok = AutoTokenizer.from_pretrained(MODEL_ID, use_fast=False)
        print("Tokenizer loaded (slow fallback).", flush=True)

    model = AutoModelForCausalLM.from_pretrained(
        MODEL_ID, torch_dtype=torch.float16, device_map="auto",
    ).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)

    df = pd.read_csv("/tmp/data/test.csv", dtype=str).fillna("")
except Exception as e:
    emergency_submission_csv(f"tokenizer/model load or test.csv read failed: {e}")
    raise

n_rows = len(df)
# Measured, not estimated: if real loading took longer or shorter than the
# SETUP_BUFFER_S guess, every downstream timing decision for the rest of the
# run should work from what actually happened, not a stale assumption.
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 (estimated {SETUP_BUFFER_S}s) | "
      f"per_row_budget={per_row_budget:.0f}s for {n_rows} rows", flush=True)


# ---- 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."


# =============================================================================
# SYMBOLIC PREPROCESSING LAYER -- pure Python standard library only (re,
# difflib, collections), no new dependencies. Deterministic, CPU-only,
# negligible runtime (contexts have ~10-20 short strings; all comparisons
# are microseconds). Survived a multi-round falsification pass: only the
# two evidence objects that (a) compute something a fast read is likely to
# miss by construction and (b) cannot mislead when wrong (worst case is
# silence, never false confidence) were kept. Augments the raw context;
# never replaces or rewrites any of it.
# =============================================================================
from difflib import SequenceMatcher
from collections import defaultdict


def extract_forms_from_context(context: str):
    """Pulls candidate unknown-language 'forms' out of raw context text, for
    reduplication's per-word self-check ONLY. Pipe-delimited lines
    contribute ONLY their FIRST field (the conventional unknown-language
    side) -- NOT every field, because including gloss/meaning fields lets
    ordinary English words (e.g. 'banana') trigger false reduplication
    hits. Plain lines contribute whitespace tokens. Lines with more than 3
    pipes are skipped defensively -- Hadza (a confirmed IOL 2026 language)
    is a click language, and '|' is sometimes used informally to transcribe
    click consonants, which would misparse as our field delimiter."""
    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):
    """Extracts genuine (input, output) pairs from pipe-delimited rows --
    e.g. fill_blanks' 'given | derived | gloss' structure -- using the row's
    own layout, not language-specific assumptions. This is the ONLY source
    of pairs fed to transformation-family detection: forms from DIFFERENT
    rows are never cross-compared, which would otherwise manufacture
    spurious 'transformations' between unrelated words. Lines with more
    than 3 pipes are skipped (see extract_forms_from_context)."""
    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):
    """A clean single-region transformation signature between two strings,
    or None if the difference is scattered across multiple regions (too
    noisy to call one transformation), OR if there is no genuine shared
    stem of at least 2 characters -- without this check, two totally
    unrelated words with zero characters in common (e.g. 'xyz' vs 'qrs')
    were being accepted as a fake 'prefix change' signature, since
    SequenceMatcher returns a single 'replace' opcode for a total mismatch
    just as it does for a real, small, genuine edit."""
    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):
    """Clusters GENUINELY PAIRED forms (same row only) sharing an identical
    clean edit signature. Emits a family only if 2+ separate given pairs
    share it -- a single occurrence is indistinguishable from coincidence
    and is worse than silence."""
    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])  # strongest support first
    return [f for _, f in families]


def detect_reduplication(forms):
    """Flags a word only if it contains an exact adjacent doubled substring
    (length >= 2). Emits nothing if absent."""
    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:
    """The full symbolic layer. Returns "" if no supported transformation
    family and no reduplication is found -- augments the prompt only when
    it has real, multi-supported evidence to add. Never replaces context."""
    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):
    """The frozen single-call decomposition scaffold -- this is the actual
    architecture behind the accepted 0.083/0.0296/0.2323 baseline (two-stage
    was tested separately and scored worse, so it is not 'current' and is
    not what this experiment augments). ONLY CHANGE: one new line -- the
    symbolic evidence block, inserted between raw context and everything
    else, per the required prompt shape (Raw Context + SYMBOLIC EVIDENCE +
    Original Question). Nothing else in this function differs from the
    frozen version: same decomposition slots, same task guidance, same
    COMPUTE note, same output contract."""
    preamble, items, count_known = parse_items(query)
    guidance = TASK_GUIDANCE.get(task_type, DEFAULT_GUIDANCE)
    symbolic_evidence = build_symbolic_evidence(context)  # "" if nothing found

    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()}"
            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}"
        )
    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}"
        )
    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:
    # Broadened: strips "Answer N:", "is:", "the answer is:", "final answer:"
    # -- Problem 1's "is: uolms" was an exact-match near-miss lost to exactly
    # this kind of un-stripped prefix.
    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(" .\"'\u201c\u201d\u2018\u2019")


def extract(text):
    """Fixed against three 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;
    (3) NO marker found + all answers dumped on one pipe-separated line --
        the fallback used to return that whole line as ONE answer instead
        of splitting it, collapsing e.g. 6 real answers into 1 giant string."""
    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()
    # Fallback (no marker found): split each line further by "|" if present,
    # instead of treating a whole pipe-separated line as one answer.
    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: 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 Exception:
        # Broadened to bare Exception: a missing chat_template raises
        # ValueError, API mismatches can be AttributeError or TypeError, and
        # a Jinja2 templating error inherits from none of those. Given the
        # stated priority is guaranteed execution, catch anything here and
        # fall back to the simpler non-dict pattern.
        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()

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

            # 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: same frozen mechanism as the baseline (dedicated
            # call if time is comfortable, else cheap truncated fallback),
            # built from the single call's own output text.
            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"])

            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)

except Exception as e:
    # Final safety net: even if something escapes every inner try/except
    # above, whatever rows were collected so far still get written.
    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)