EjZhou commited on
Commit
a4583db
·
verified ·
1 Parent(s): 3a707f8

Upload script.py with huggingface_hub

Browse files
Files changed (1) hide show
  1. script.py +546 -0
script.py ADDED
@@ -0,0 +1,546 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ IOL-AI Challenge 2026 — submission script (OFFLINE / Mode B).
3
+
4
+ Runtime facts (Space Submission tab):
5
+ * T4 medium, 16 GB VRAM, Python 3.10, 30-min wall clock.
6
+ * NO internet: cannot pip install or download anything. Model weights must be
7
+ committed into THIS repo (the working dir) and loaded from ".". Only the
8
+ pre-installed libraries/versions are available (torch 2.4.0, transformers
9
+ 4.44.1, accelerate 0.34.2, bitsandbytes 0.43.3, autoawq 0.2.7, pandas 2.2.2,
10
+ numpy 2.1.3, ...). Do NOT pin different majors of torch/transformers/numpy.
11
+ * Read hidden test set from /tmp/data/test.csv; write submission.csv here.
12
+ * pred = JSON list, one entry per numbered item, in query order.
13
+
14
+ Ship the model in the repo with build_repo.py. This script loads it from "." with
15
+ bitsandbytes 4-bit by default (or auto-detected AWQ) so it fits 16 GB. T4 has no
16
+ bf16 -> use float16.
17
+
18
+ Lineage: v4 fixed the big bug — answer COUNT comes from the model/context, never a
19
+ query regex (that clipped matching/fill-blank problems to 1 answer). v10 added
20
+ deterministic exact-match boosters (arithmetic-eval for numbers, bijection repair for
21
+ matching, diacritic/gloss fidelity). v11 (this file) spends the whole 30-min wall
22
+ SAFELY: per row, a greedy anchor + sampled self-consistency (majority vote, early stop
23
+ on consensus) + an optional refine pass ("re-derive, check, fix", folded in as one more
24
+ vote) — all bounded by an adaptive per-row budget + per-decode max_time + soft/hard wall
25
+ phases that degrade to single-pass then placeholder, so it can never time out (v6 did).
26
+ Single-sequence decode throughout (v3's batched decode OOM'd). Every row keeps an
27
+ `explanation` (Human-Eval track). Tunable: IOL_MAX_SAMPLES / IOL_CONSENSUS / IOL_REFINE
28
+ / IOL_TEMPERATURE / IOL_TOP_P / IOL_MAX_NEW_TOKENS / IOL_SOFT_BUDGET_S / IOL_HARD_BUDGET_S.
29
+
30
+ Local dev: set IOL_TEST_CSV to a mock file. Quantization auto-disables if there's
31
+ no CUDA so the plumbing can be exercised on CPU with a tiny model.
32
+ """
33
+
34
+ import os
35
+ os.environ.setdefault("HF_HUB_OFFLINE", "1")
36
+ os.environ.setdefault("TRANSFORMERS_OFFLINE", "1")
37
+
38
+ import re
39
+ import csv
40
+ import json
41
+ import time
42
+
43
+ SCRIPT_START = time.time() # ~wall-clock start; the 30-min hard kill counts from here
44
+
45
+ MODEL_DIR = os.environ.get("IOL_MODEL_DIR", ".") # weights live in the repo
46
+ TEST_CSV = os.environ.get("IOL_TEST_CSV", "/tmp/data/test.csv")
47
+ OUT_CSV = os.environ.get("IOL_OUT_CSV", "submission.csv")
48
+ MAX_NEW_TOKENS = int(os.environ.get("IOL_MAX_NEW_TOKENS", "768"))
49
+ # "4bit" (bitsandbytes), "awq" (weights already AWQ-quantized), or "fp16".
50
+ QUANT = os.environ.get("IOL_QUANT", "4bit")
51
+
52
+ # --- v11: adaptive compute — self-consistency + refine, bounded by the wall -----
53
+ # Per row: a greedy ANCHOR decode + extra SAMPLED decodes, majority-voted per item
54
+ # (self-consistency), stopping early on consensus so easy rows free up time for hard
55
+ # ones; then, budget permitting, a REFINE pass re-checks the draft against the data and
56
+ # is folded in as one more vote. Everything is gated by an ADAPTIVE per-row time budget
57
+ # and a per-decode max_time, so we spend the whole wall but ALWAYS finish (v6 timed out
58
+ # for lack of exactly this). Single-sequence decode throughout (no OOM).
59
+ TEMPERATURE = float(os.environ.get("IOL_TEMPERATURE", "0.7"))
60
+ TOP_P = float(os.environ.get("IOL_TOP_P", "0.9"))
61
+ MAX_SAMPLES = int(os.environ.get("IOL_MAX_SAMPLES", "6")) # hard cap of decodes/row
62
+ CONSENSUS_FRAC = float(os.environ.get("IOL_CONSENSUS", "0.6")) # early-stop agreement level
63
+ DO_REFINE = os.environ.get("IOL_REFINE", "1") != "0"
64
+ MIN_DECODE_S = float(os.environ.get("IOL_MIN_DECODE_S", "20")) # floor for a decode's max_time
65
+ # Wall-clock phases (seconds from SCRIPT_START; 30-min hard kill = 1800):
66
+ SOFT_BUDGET_S = float(os.environ.get("IOL_SOFT_BUDGET_S", "1560")) # 26.0m: full pipeline before this
67
+ HARD_BUDGET_S = float(os.environ.get("IOL_HARD_BUDGET_S", "1710")) # 28.5m: placeholder-fill after
68
+
69
+ ANSWER_MARKER = "###ANSWERS###"
70
+ WHY_MARKER = "###WHY###"
71
+
72
+ SYSTEM_PROMPT = (
73
+ "You are an expert competitor at the International Linguistics Olympiad. "
74
+ "Each problem gives data from a language you have never seen; deduce its rules "
75
+ "using ONLY the data and hints in the problem, then answer EVERY sub-question.\n\n"
76
+ "A problem can have MANY sub-questions even when the query is one sentence: e.g. "
77
+ "'give the correspondences' expects one answer for EACH numbered item in the data "
78
+ "(often a dozen or more). Work out how many answers are required and give exactly "
79
+ "that many, one per item, in the order the items appear.\n\n"
80
+ "Reason briefly, then give your answers in EXACTLY this format:\n"
81
+ f"{ANSWER_MARKER}\n"
82
+ "1. <answer to item 1>\n"
83
+ "2. <answer to item 2>\n"
84
+ "(one numbered line per sub-question, in order)\n"
85
+ f"{WHY_MARKER}\n"
86
+ "- <the key rule or pattern you found>\n"
87
+ "- <the main evidence from the data that supports it>\n\n"
88
+ "Each answer line holds ONLY the requested form — a word, phrase, number, or "
89
+ "letter — with no restating of the question and no commentary. Answer in the "
90
+ "language and direction the query asks. For matching items give just the option "
91
+ "letter; for number items give digits or the written-out number as asked. Never "
92
+ "leave an item blank — always give your best guess.\n\n"
93
+ "EXACT SPELLING MATTERS: copy the exact characters, diacritics and special symbols "
94
+ "that appear in the data (e.g. ʼ ɨ ŋ ʂ); never swap them for similar-looking "
95
+ "ordinary letters. When translating INTO English, reproduce the examples' glossing "
96
+ "style verbatim, including person/number markers written like you_sg, you_pl.\n\n"
97
+ f"The lines after {WHY_MARKER} are a SHORT, human-readable explanation (1-3 bullets "
98
+ "a person can grasp in under a minute) — NOT your full reasoning trace."
99
+ )
100
+
101
+ # Short, low-cost per-task output reminders (the CSV tags each row with task_type).
102
+ TASK_HINT = {
103
+ "translation": "This is a translation task: each answer is only the translated word/phrase.",
104
+ "text_to_num": "This is a number task: each answer is only digits (e.g. 42).",
105
+ "num_to_text": "This is a number task: each answer is only the number written in the target language's words.",
106
+ "match_letters": "This is a matching task: each answer is only the option letter (A, B, C, ...); give one per item in the data.",
107
+ "matching": "This is a matching task: each answer is only the option letter; give one per item in the data.",
108
+ "fill_blank": "This is a fill-in-the-blank task: each answer is only the missing form.",
109
+ "fill_blanks": "This is a fill-in-the-blank task: each answer is only the missing form.",
110
+ }
111
+
112
+
113
+ def build_messages(row):
114
+ """Chat messages for one problem, with a short task_type-specific reminder plus
115
+ two exact-match boosters: a COMPUTE line for number tasks (we evaluate it) and the
116
+ exact set of valid letters for matching tasks."""
117
+ context = (row.get("context") or "").strip()
118
+ query = (row.get("query") or "").strip()
119
+ ttype = (row.get("task_type") or "").strip().lower()
120
+
121
+ system = SYSTEM_PROMPT
122
+ hint = TASK_HINT.get(ttype)
123
+ if hint:
124
+ system = system + "\n\n" + hint
125
+
126
+ if ttype == "text_to_num":
127
+ system += (
128
+ f"\n\nAfter the {WHY_MARKER} bullets, add one more line exactly:\n"
129
+ "COMPUTE: expr1 | expr2 | ...\n"
130
+ "where each expr is plain arithmetic (digits, + - *, parentheses only) that "
131
+ "evaluates to that item's number, one per item, matching the rule you found."
132
+ )
133
+ if ttype in ("match_letters", "matching"):
134
+ opts = extract_letter_options(context)
135
+ if opts:
136
+ system += (f"\n\nThe ONLY valid answers are these letters: {', '.join(opts)}. "
137
+ "Use no other letter.")
138
+
139
+ return [
140
+ {"role": "system", "content": system},
141
+ {"role": "user", "content": context + "\n\n" + query},
142
+ ]
143
+
144
+
145
+ def detect_count(context, query):
146
+ """Best-effort number of sub-questions — used ONLY as a hint and a minimum pad,
147
+ NEVER to truncate the model's own answer list (under-producing loses items).
148
+ Queries usually number items ('1.'/'2.') or mark blanks ('(1)','(2)'); matching
149
+ queries number nothing, so fall back to the numbered items in the CONTEXT."""
150
+ q = re.findall(r"(?m)^\s*(\d+)[\.\)]", query)
151
+ if q:
152
+ return len(q)
153
+ par = re.findall(r"\((\d+)\)", query)
154
+ if par:
155
+ return len(set(par))
156
+ c = re.findall(r"(?m)^\s*(\d+)[\.\)]", context)
157
+ if c:
158
+ return len(c)
159
+ return 1
160
+
161
+
162
+ def _clean_answer(s):
163
+ """Strip list markers, common 'Answer:' labels, and surrounding quotes."""
164
+ s = re.sub(r"^\s*(?:\d+[\.\):]|[-*•])\s*", "", s).strip()
165
+ s = re.sub(r"^(?:answer|ans|translation|result)\s*[:\-]\s*", "", s, flags=re.I).strip()
166
+ return s.strip("\"'“”‘’` ").strip()
167
+
168
+
169
+ def parse_answers(text, min_count=1):
170
+ """Extract the model's FULL answer list — the count comes from the MODEL, never
171
+ truncated to a query heuristic (that was the v1/v2 bug: it clipped matching
172
+ problems' dozen answers down to 1). Prefer the ###ANSWERS### block; inside it
173
+ read the numbered lines; else split a trailing comma-list; else use the lines.
174
+ Pad up to min_count and never emit a blank."""
175
+ seg = text.rsplit(ANSWER_MARKER, 1)[1] if ANSWER_MARKER in text else text
176
+ if WHY_MARKER in seg: # answers live BEFORE the WHY section
177
+ seg = seg.split(WHY_MARKER, 1)[0]
178
+
179
+ numbered = {}
180
+ for m in re.finditer(r"(?m)^\s*(\d+)[\.\)]\s*(.+?)\s*$", seg):
181
+ numbered[int(m.group(1))] = _clean_answer(m.group(2))
182
+ if numbered: # ordered by the model's indices
183
+ answers = [numbered.get(i, "") for i in range(1, max(numbered) + 1)]
184
+ else:
185
+ lines = [ln.strip() for ln in seg.splitlines() if ln.strip()]
186
+ comma_line = next((ln for ln in reversed(lines) if "," in ln), "")
187
+ if comma_line: # matching-style "O, D, A, ..."
188
+ answers = [_clean_answer(x) for x in comma_line.split(",")]
189
+ else:
190
+ answers = [_clean_answer(ln) for ln in lines]
191
+
192
+ answers = [a if a else "?" for a in answers] # never blank (partial credit)
193
+ if len(answers) < min_count:
194
+ answers += ["?"] * (min_count - len(answers))
195
+ return answers if answers else ["?"]
196
+
197
+
198
+ def _norm(s):
199
+ """Mirror the official scorer's normalization so voting groups answers the
200
+ same way the metric will (ignore case, surrounding quotes, one trailing dot)."""
201
+ s = " ".join((s or "").strip().split())
202
+ s = s.strip("\"'“”‘’")
203
+ if s.endswith("."):
204
+ s = s[:-1]
205
+ return s.strip().casefold()
206
+
207
+
208
+ def vote_lists(answer_lists, min_count=1):
209
+ """Majority-vote per position across already-parsed answer lists (self-consistency).
210
+ Most common NORMALIZED form wins; returns its surface form. Ties -> earliest list.
211
+ '?' placeholders don't vote, so a malformed/truncated decode can't corrupt a result."""
212
+ from collections import Counter
213
+
214
+ n = max([min_count] + [len(a) for a in answer_lists]) if answer_lists else min_count
215
+ out = []
216
+ for i in range(n):
217
+ counts, surface = Counter(), {}
218
+ for a in answer_lists:
219
+ if i < len(a) and a[i] and a[i] != "?":
220
+ key = _norm(a[i])
221
+ counts[key] += 1
222
+ surface.setdefault(key, a[i])
223
+ out.append(surface[counts.most_common(1)[0][0]] if counts else "?")
224
+ return out
225
+
226
+
227
+ def vote_answers(sample_texts, min_count=1):
228
+ """Self-consistency over raw decodes (parse each, then vote per item)."""
229
+ return vote_lists([parse_answers(t, min_count) for t in sample_texts], min_count)
230
+
231
+
232
+ def _agreement(answer_lists, voted):
233
+ """Fraction of lists whose normalized answers equal the voted result (early-stop)."""
234
+ if not answer_lists:
235
+ return 0.0
236
+ vt = tuple(_norm(x) for x in voted)
237
+ hit = 0
238
+ for a in answer_lists:
239
+ ap = list(a) + ["?"] * (len(vt) - len(a))
240
+ if tuple(_norm(x) for x in ap[:len(vt)]) == vt:
241
+ hit += 1
242
+ return hit / len(answer_lists)
243
+
244
+
245
+ def parse_explanation(text):
246
+ """Pull the short ###WHY### summary the model wrote (for the Human-Eval jury track).
247
+ Kept concise and readable; NOT the raw reasoning trace. '' if the model omitted it.
248
+ The internal COMPUTE: line (used only for arithmetic eval) is dropped from it."""
249
+ if WHY_MARKER not in text:
250
+ return ""
251
+ why = text.rsplit(WHY_MARKER, 1)[1].replace(ANSWER_MARKER, " ").strip()
252
+ lines = [ln.strip() for ln in why.splitlines()
253
+ if ln.strip() and not re.match(r"(?i)^\s*compute\s*:", ln)]
254
+ return "\n".join(lines[:4])[:600].strip()
255
+
256
+
257
+ # ---- deterministic exact-match boosters (no extra model call, adapted from v5) ----
258
+ import ast as _ast
259
+ _ALLOWED_BINOPS = (_ast.Add, _ast.Sub, _ast.Mult)
260
+ _NUM_NODE = getattr(_ast, "Num", None) # pre-3.8 number node (sandbox is 3.10=Constant)
261
+
262
+
263
+ def _safe_arithmetic(expr):
264
+ """Evaluate a plain +,-,* / parenthesised integer expression, else None."""
265
+ try:
266
+ tree = _ast.parse(expr.strip(), mode="eval")
267
+ except Exception:
268
+ return None
269
+
270
+ def _ev(n):
271
+ if isinstance(n, _ast.Expression):
272
+ return _ev(n.body)
273
+ if isinstance(n, _ast.Constant) and isinstance(n.value, (int, float)):
274
+ return n.value
275
+ if _NUM_NODE is not None and isinstance(n, _NUM_NODE): # Python <3.8
276
+ return n.n
277
+ if isinstance(n, _ast.BinOp) and isinstance(n.op, _ALLOWED_BINOPS):
278
+ l, r = _ev(n.left), _ev(n.right)
279
+ if l is None or r is None:
280
+ return None
281
+ if isinstance(n.op, _ast.Add):
282
+ return l + r
283
+ if isinstance(n.op, _ast.Sub):
284
+ return l - r
285
+ return l * r
286
+ if isinstance(n, _ast.UnaryOp) and isinstance(n.op, _ast.USub):
287
+ v = _ev(n.operand)
288
+ return -v if v is not None else None
289
+ return None
290
+
291
+ return _ev(tree)
292
+
293
+
294
+ def apply_compute_overrides(text, answers):
295
+ """text_to_num: if the model wrote 'COMPUTE: e1 | e2', evaluate each safely and
296
+ override that item's answer with the exact integer — kills arithmetic slips while
297
+ keeping the model's derived rule. Only overrides when the eval is a clean integer."""
298
+ m = re.search(r"(?im)^\s*COMPUTE\s*:\s*(.+)$", text)
299
+ if not m:
300
+ return answers
301
+ exprs = [e.strip() for e in m.group(1).split("|")]
302
+ out = list(answers)
303
+ for i, e in enumerate(exprs[:len(out)]):
304
+ v = _safe_arithmetic(e)
305
+ if v is not None and float(v).is_integer():
306
+ out[i] = str(int(v))
307
+ return out
308
+
309
+
310
+ def extract_letter_options(context):
311
+ """The option labels A, B, C ... that a matching problem offers (contiguous from A)."""
312
+ found = set()
313
+ for line in context.splitlines():
314
+ for m in re.finditer(r"(?:^|\s)([A-Z])[.\)]\s+\S", line):
315
+ found.add(m.group(1))
316
+ if not found:
317
+ return None
318
+ letters = sorted(found)
319
+ if letters != [chr(ord("A") + i) for i in range(len(letters))]:
320
+ return None
321
+ return letters if 2 <= len(letters) <= 26 else None
322
+
323
+
324
+ def repair_bijection(answers, labels):
325
+ """match_letters, bijection case only: keep the letters the model committed to
326
+ (first occurrence wins), fill duplicate/invalid/missing slots with the leftover
327
+ letters in order -> a guaranteed valid permutation of exactly len(labels) items."""
328
+ n = len(labels)
329
+ labels_sorted = sorted(labels)
330
+ picks = []
331
+ for i in range(n):
332
+ a = answers[i] if i < len(answers) else ""
333
+ f = re.findall(r"[A-Za-z]", a or "")
334
+ c = f[0].upper() if f else ""
335
+ picks.append(c if c in labels else "")
336
+ result = [None] * n
337
+ used = set()
338
+ for i in range(n):
339
+ if picks[i] and picks[i] not in used:
340
+ result[i] = picks[i]
341
+ used.add(picks[i])
342
+ missing = [l for l in labels_sorted if l not in used]
343
+ mi = 0
344
+ for i in range(n):
345
+ if result[i] is None:
346
+ result[i] = missing[mi] if mi < len(missing) else labels_sorted[0]
347
+ mi += 1
348
+ return result
349
+
350
+
351
+ def postprocess(row, answers, raw):
352
+ """Apply the deterministic boosters that fit this row's task_type."""
353
+ ttype = (row.get("task_type") or "").strip().lower()
354
+ context = (row.get("context") or "")
355
+ if ttype == "text_to_num":
356
+ answers = apply_compute_overrides(raw, answers)
357
+ elif ttype in ("match_letters", "matching"):
358
+ labels = extract_letter_options(context)
359
+ # Only a true bijection (numbered context items == number of labels) is safe to
360
+ # repair; "pick the letter for item 1,2" style (few items, reused letters) is not.
361
+ ctx_items = len(re.findall(r"(?m)^\s*\d+\s*[.\)]", context))
362
+ if labels and len(labels) >= 2 and ctx_items == len(labels):
363
+ answers = repair_bijection(answers, labels)
364
+ return answers
365
+
366
+
367
+ def default_explanation(row):
368
+ """Non-empty fallback so the explanation column is populated on every row."""
369
+ t = (row.get("task_type") or "linguistic").replace("_", " ")
370
+ return f"Inferred the {t} rule from the given examples and applied it to each item."
371
+
372
+
373
+ def _already_quantized(model_dir):
374
+ """True if the shipped weights are pre-quantized (e.g. AWQ) — then transformers
375
+ auto-detects the config and we must NOT stack bitsandbytes on top."""
376
+ cfg = os.path.join(model_dir, "config.json")
377
+ try:
378
+ with open(cfg, encoding="utf-8") as f:
379
+ return "quantization_config" in json.load(f)
380
+ except Exception:
381
+ return False
382
+
383
+
384
+ def load_model():
385
+ import torch
386
+ from transformers import AutoTokenizer, AutoModelForCausalLM
387
+
388
+ tok = AutoTokenizer.from_pretrained(MODEL_DIR)
389
+ if tok.pad_token_id is None: # only used to silence a warning
390
+ tok.pad_token = tok.eos_token
391
+ if not torch.cuda.is_available():
392
+ model = AutoModelForCausalLM.from_pretrained(
393
+ MODEL_DIR, torch_dtype=torch.float32).eval() # CPU dev fallback
394
+ return tok, model
395
+
396
+ kwargs = dict(torch_dtype=torch.float16, device_map="auto") # T4 has no bf16
397
+ if _already_quantized(MODEL_DIR):
398
+ pass # AWQ/pre-quant: transformers reads quantization_config from config.json
399
+ elif QUANT == "4bit":
400
+ from transformers import BitsAndBytesConfig
401
+ kwargs["quantization_config"] = BitsAndBytesConfig(
402
+ load_in_4bit=True,
403
+ bnb_4bit_compute_dtype=torch.float16,
404
+ bnb_4bit_quant_type="nf4",
405
+ bnb_4bit_use_double_quant=True,
406
+ )
407
+ model = AutoModelForCausalLM.from_pretrained(MODEL_DIR, **kwargs).eval()
408
+ return tok, model
409
+
410
+
411
+ def generate_one(tok, model, messages, do_sample, max_time=None):
412
+ """Single-sequence decode (batch=1) — the VRAM-safe path proven in v1/v2. (v3's
413
+ batched multi-sequence decode OOM'd the T4 and produced an empty submission.)
414
+ max_time caps this one decode's wall time so the pipeline can't overrun."""
415
+ import torch
416
+
417
+ dev = model.device if hasattr(model, "device") else "cpu"
418
+ ids = tok.apply_chat_template(
419
+ messages, add_generation_prompt=True, return_tensors="pt").to(dev)
420
+ gkw = dict(max_new_tokens=MAX_NEW_TOKENS, pad_token_id=tok.pad_token_id)
421
+ if do_sample:
422
+ gkw.update(do_sample=True, temperature=TEMPERATURE, top_p=TOP_P)
423
+ else:
424
+ gkw.update(do_sample=False)
425
+ if max_time and max_time > 0:
426
+ gkw["max_time"] = float(max_time) # MaxTimeCriteria stops the decode
427
+ with torch.no_grad():
428
+ gen = model.generate(ids, **gkw)
429
+ return tok.decode(gen[0][ids.shape[-1]:], skip_special_tokens=True).strip()
430
+
431
+
432
+ def build_refine_messages(row, answers):
433
+ """A 'check and correct' pass: show the draft answers and ask the model to re-derive
434
+ the rule from the data, verify each answer, and output a corrected list (same format)."""
435
+ context = (row.get("context") or "").strip()
436
+ query = (row.get("query") or "").strip()
437
+ draft = "\n".join("%d. %s" % (i + 1, a) for i, a in enumerate(answers))
438
+ system = (
439
+ "You are re-checking a DRAFT solution to an International Linguistics Olympiad "
440
+ "problem. Re-derive the rules STRICTLY from the data, test them on every example, "
441
+ "then verify each draft answer and FIX any that are wrong (keep the right ones). "
442
+ "Output the corrected answers in EXACTLY this format:\n"
443
+ f"{ANSWER_MARKER}\n1. <answer 1>\n2. <answer 2>\n(one per sub-question, in order)\n"
444
+ f"{WHY_MARKER}\n- <the rule you used>\n- <what you changed, if anything>\n\n"
445
+ "Copy exact characters/diacritics; give only the requested form; never leave a blank."
446
+ )
447
+ user = "%s\n\n%s\n\nDRAFT ANSWERS:\n%s" % (context, query, draft)
448
+ return [{"role": "system", "content": system}, {"role": "user", "content": user}]
449
+
450
+
451
+ def _well_formed(raw, answer_list, min_count):
452
+ """A refine/sample result is trustworthy enough to vote only if it hit the marker
453
+ and yielded at least min_count real (non-'?') answers."""
454
+ real = sum(1 for a in answer_list if a and a != "?")
455
+ return (ANSWER_MARKER in raw) and real >= min_count
456
+
457
+
458
+ def solve_row(tok, model, row, row_deadline):
459
+ """Adaptive per-row compute: greedy anchor + sampled self-consistency (early stop on
460
+ consensus) + an optional refine vote, all bounded by row_deadline and the wall phases.
461
+ Returns (answers, explanation, mode)."""
462
+ context = (row.get("context") or "").strip()
463
+ query = (row.get("query") or "").strip()
464
+ min_count = detect_count(context, query)
465
+ messages = build_messages(row)
466
+
467
+ def budget(): # time left for THIS row's next decode
468
+ by_row = row_deadline - time.time()
469
+ by_wall = SOFT_BUDGET_S - (time.time() - SCRIPT_START)
470
+ return min(by_row, by_wall)
471
+
472
+ elapsed = time.time() - SCRIPT_START
473
+ if elapsed > HARD_BUDGET_S: # no time: safe placeholder, keep file whole
474
+ return ["?"] * min_count, default_explanation(row), "placeholder"
475
+ if elapsed > SOFT_BUDGET_S: # low time: one greedy pass only
476
+ raw = generate_one(tok, model, messages, False, max_time=max(MIN_DECODE_S, budget()))
477
+ return (postprocess(row, parse_answers(raw, min_count), raw),
478
+ parse_explanation(raw) or default_explanation(row), "single")
479
+
480
+ # --- full pipeline: greedy anchor, then sampled self-consistency ---
481
+ anchor = generate_one(tok, model, messages, False, max_time=max(MIN_DECODE_S, budget()))
482
+ texts = [anchor]
483
+ lists = [postprocess(row, parse_answers(anchor, min_count), anchor)]
484
+ while len(texts) < MAX_SAMPLES and budget() > MIN_DECODE_S:
485
+ t = generate_one(tok, model, messages, True, max_time=budget())
486
+ texts.append(t)
487
+ lists.append(postprocess(row, parse_answers(t, min_count), t))
488
+ if len(lists) >= 3 and _agreement(lists, vote_lists(lists, min_count)) >= CONSENSUS_FRAC:
489
+ break
490
+
491
+ voted = vote_lists(lists, min_count)
492
+ explanation = parse_explanation(anchor) or default_explanation(row)
493
+ mode = "vote%d" % len(texts)
494
+
495
+ # --- optional refine pass, folded in as one more vote (only if well-formed) ---
496
+ if DO_REFINE and budget() > MIN_DECODE_S:
497
+ rtext = generate_one(tok, model, build_refine_messages(row, voted), False,
498
+ max_time=budget())
499
+ rlist = postprocess(row, parse_answers(rtext, min_count), rtext)
500
+ if _well_formed(rtext, rlist, min_count):
501
+ voted = vote_lists(lists + [rlist], min_count)
502
+ explanation = parse_explanation(rtext) or explanation
503
+ mode += "+refine"
504
+
505
+ return voted, explanation, mode
506
+
507
+
508
+ def main():
509
+ tok, model = load_model()
510
+
511
+ with open(TEST_CSV, newline="", encoding="utf-8") as f:
512
+ rows = list(csv.DictReader(f))
513
+
514
+ # Write incrementally so a hard 30-min kill still leaves a valid partial file.
515
+ # 'explanation' column opts into the Human-Eval jury track (not auto-scored).
516
+ fout = open(OUT_CSV, "w", newline="", encoding="utf-8")
517
+ writer = csv.DictWriter(fout, fieldnames=["id", "pred", "explanation"])
518
+ writer.writeheader()
519
+ fout.flush()
520
+
521
+ n = len(rows)
522
+ for k, r in enumerate(rows):
523
+ rows_left = n - k
524
+ soft_left = SOFT_BUDGET_S - (time.time() - SCRIPT_START)
525
+ # Even share of the remaining soft budget; rows that finish early hand their
526
+ # slack to later rows (next iteration's soft_left is larger).
527
+ row_deadline = time.time() + max(0.0, soft_left) / max(1, rows_left)
528
+ try:
529
+ answers, explanation, mode = solve_row(tok, model, r, row_deadline)
530
+ except Exception as e: # one row must never zero the whole submission
531
+ print("row %s failed: %r" % (r.get("id"), e), flush=True)
532
+ mc = detect_count((r.get("context") or ""), (r.get("query") or ""))
533
+ answers, explanation, mode = ["?"] * mc, default_explanation(r), "error"
534
+ writer.writerow({"id": r["id"],
535
+ "pred": json.dumps(answers, ensure_ascii=False),
536
+ "explanation": explanation})
537
+ fout.flush() # survive a hard timeout
538
+ print("%d/%d [%s] elapsed=%ds" % (k + 1, n, mode, time.time() - SCRIPT_START),
539
+ flush=True)
540
+
541
+ fout.close()
542
+ print("wrote %s (%d rows)" % (OUT_CSV, n), flush=True)
543
+
544
+
545
+ if __name__ == "__main__":
546
+ main()