arvindcr4 commited on
Commit
24d7d8e
·
verified ·
1 Parent(s): 5de2079

v12: assignment solver for match_letters; single change on the 0.2245 v8 config

Browse files
Files changed (1) hide show
  1. script.py +817 -0
script.py ADDED
@@ -0,0 +1,817 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python
2
+ """IOL-AI 2026 submission -- International Linguistics Olympiad solver.
3
+
4
+ Design notes (the eval sandbox is unforgiving, so these matter):
5
+
6
+ * HARD 30-MINUTE LIMIT. A killed process means no score at all, so the script
7
+ is structured as a monotonically-improving pipeline: it writes a complete,
8
+ correctly-shaped submission.csv *before* the model is even loaded, then
9
+ overwrites it after every improvement. Any crash or timeout leaves the best
10
+ result reached so far on disk.
11
+ * ALIGNMENT IS EVERYTHING. Each row is a problem block with N numbered items
12
+ and `pred` must be a JSON list of exactly N answers, in order. One missing
13
+ line shifts every later answer and zeroes the whole block on both metrics.
14
+ So N is detected from the query and the model output is force-fitted to it.
15
+ * NEVER EMIT AN EMPTY STRING. The final score is a geometric mean of exact
16
+ match and chrF, so an empty answer scores zero on both. A wrong guess is
17
+ strictly better than a blank.
18
+ * Environment is transformers 4.44.1 / torch 2.4.0 / autoawq on a 16GB T4
19
+ (fp16 only, no bf16, no flash-attn), with no internet.
20
+ """
21
+ import os
22
+ import re
23
+ import json
24
+ import time
25
+ import unicodedata
26
+ from collections import Counter, defaultdict
27
+
28
+ T0 = time.time()
29
+
30
+ # The platform allows 30 minutes. Reserve a margin for model load overhead we
31
+ # can't predict and for the final write; being 60s early costs a little
32
+ # accuracy, being 1s late costs the entire submission.
33
+ TIME_LIMIT = float(os.environ.get("IOL_TIME_LIMIT", "1800"))
34
+ SAFETY = float(os.environ.get("IOL_SAFETY", "150"))
35
+ DEADLINE = T0 + TIME_LIMIT - SAFETY
36
+
37
+ TEST_CSV = os.environ.get("IOL_TEST_CSV", "/tmp/data/test.csv")
38
+ OUT_CSV = os.environ.get("IOL_OUT_CSV", "submission.csv")
39
+ MODEL_ID = os.environ.get("IOL_MODEL", ".")
40
+ WANT_EXPLANATION = os.environ.get("IOL_EXPLAIN", "1") == "1"
41
+ MAX_NEW = int(os.environ.get("IOL_MAXNEW", "900")) # reasoning budget/item
42
+ MAX_SAMPLES = int(os.environ.get("IOL_MAXSAMPLES", "8")) # self-consistency cap
43
+ # BASELINE REPLICATION MODE. The organizers' reference script reaches exact match
44
+ # 0.0729 on the hidden set with THESE EXACT WEIGHTS; our best is 0.0333. Before
45
+ # adding anything else we need to know whether that number is reproducible by us
46
+ # at all. This mode replicates their script literally -- trivial prompt, no CoT,
47
+ # 512 tokens, batch 1 (no padding at all), naive line split, NO forcing to N --
48
+ # and changes exactly one thing: repetition_penalty=1.0, our one proven fix.
49
+ BASELINE_MODE = os.environ.get("IOL_BASELINE", "1") == "1" # v8: ON by default
50
+ # Lower than the usual 0.7: samples only earn a vote by agreeing with each
51
+ # other, so keeping them near the greedy mode makes agreement meaningful.
52
+ SAMPLE_TEMP = float(os.environ.get("IOL_TEMP", "0.5"))
53
+
54
+ os.environ.setdefault("HF_HUB_OFFLINE", "1")
55
+ os.environ.setdefault("TRANSFORMERS_OFFLINE", "1")
56
+ os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
57
+ # Reduce allocator fragmentation: at batch 4 the T4 has only ~2GB spare.
58
+ os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
59
+
60
+
61
+ def log(msg):
62
+ print(f"[{time.time() - T0:7.1f}s] {msg}", flush=True)
63
+
64
+
65
+ def left():
66
+ return DEADLINE - time.time()
67
+
68
+
69
+ # ===========================================================================
70
+ # Item-count detection (validated: 98.4% of Linguini items land in
71
+ # correctly-sized blocks)
72
+ # ===========================================================================
73
+
74
+ _LINE_NUM = re.compile(r"^[ \t]*(\d{1,3})[.)\]]", re.M)
75
+ _PAREN_NUM = re.compile(r"\((\d{1,3})\)")
76
+ _RANGE = re.compile(r"\(?(\d{1,3})\s*(?:[-–—]|to)\s*(\d{1,3})\)?")
77
+ _LINE_LETTER = re.compile(r"^[ \t]*([A-Z])[.)\]]\s", re.M)
78
+ _PAREN_LETTER = re.compile(r"\(([A-Z])\)")
79
+
80
+
81
+ def detect_n_items(query, task_type="", context=""):
82
+ """How many numbered sub-items this problem asks for. Never < 1."""
83
+ q = query or ""
84
+ line_nums = [int(m) for m in _LINE_NUM.findall(q)]
85
+ paren_nums = [int(m) for m in _PAREN_NUM.findall(q)]
86
+
87
+ range_n = 0
88
+ for a, b in _RANGE.findall(q):
89
+ a, b = int(a), int(b)
90
+ if 0 < b - a < 60:
91
+ range_n = max(range_n, b - a + 1)
92
+
93
+ cand = max(len(set(line_nums)), len(set(paren_nums)))
94
+ if range_n and cand and range_n != cand:
95
+ # A stated range ("items 1-4") can disagree with the markers actually
96
+ # present; the markers are what we have to answer, so they win.
97
+ return cand
98
+ cand = max(cand,
99
+ len(set(_LINE_LETTER.findall(q))),
100
+ len(set(_PAREN_LETTER.findall(q))))
101
+
102
+ n = max(range_n, cand)
103
+ if n > 1:
104
+ return n
105
+
106
+ # Unnumbered "Translate into X:" followed by one item per line.
107
+ lines = [l.strip() for l in q.splitlines() if l.strip()]
108
+ if len(lines) > 1:
109
+ head = lines[0]
110
+ body = lines[1:] if head.endswith((":", ".")) else lines
111
+ if body:
112
+ return len(body)
113
+
114
+ # Bare instruction ("Determine the correct correspondences."): items are in
115
+ # the shared context (this is the match_letters shape).
116
+ if context:
117
+ c_nums = len(set(int(m) for m in _LINE_NUM.findall(context)))
118
+ if c_nums > 1:
119
+ return c_nums
120
+ c_lets = len(set(_LINE_LETTER.findall(context)))
121
+ if c_lets > 1:
122
+ return c_lets
123
+
124
+ return max(n, 1)
125
+
126
+
127
+ # ===========================================================================
128
+ # Output parsing / repair
129
+ # ===========================================================================
130
+
131
+ _STRIP_PREFIX = re.compile(r"^\s*(?:\(?\d{1,3}\)?[.):\]]\s*|[-*•]\s+)")
132
+ _FENCE = re.compile(r"^```[a-zA-Z]*\s*$")
133
+ _CHATTY = re.compile(
134
+ r"^\s*(?:here (?:are|is)\b|answers?\s*:?\s*$|explanation\b|note\b|okay\b|"
135
+ r"solution\b|reasoning\b|analysis\b|translations?\s*:?\s*$|the answers?\b|"
136
+ r"let me\b|first,|so,|therefore\b|thus\b)",
137
+ re.I,
138
+ )
139
+
140
+
141
+ def clean_line(s):
142
+ s = s.strip()
143
+ s = _STRIP_PREFIX.sub("", s)
144
+ s = s.strip().strip("`").strip()
145
+ if len(s) >= 2 and s[0] == s[-1] and s[0] in "\"'“”":
146
+ s = s[1:-1].strip()
147
+ # "word | gloss" answer lines: keep the side being asked for is ambiguous,
148
+ # so keep the whole line -- chrF still gives partial credit.
149
+ return s.strip()
150
+
151
+
152
+ def extract_item_sources(query, n):
153
+ """The source text of each numbered item, used as a last-resort fallback.
154
+
155
+ A blank scores zero on both metrics; echoing the item's own source string is
156
+ strictly better, and on transcription / fill-the-blank tasks the source and
157
+ the target share a lot of characters, so it collects real chrF credit.
158
+ """
159
+ q = query or ""
160
+ out = []
161
+ for ln in q.splitlines():
162
+ s = ln.strip()
163
+ if not s:
164
+ continue
165
+ m = re.match(r"^\(?(\d{1,3})\)?[.):\]]\s*(.+)$", s)
166
+ if m:
167
+ out.append(m.group(2).strip())
168
+ if not out:
169
+ lines = [l.strip() for l in q.splitlines() if l.strip()]
170
+ if len(lines) > 1 and lines[0].endswith((":", ".")):
171
+ out = lines[1:]
172
+ # "form | gloss" items: the left side is the thing being asked about.
173
+ out = [o.split("|")[0].strip() if "|" in o else o for o in out]
174
+ out = [o for o in out if o]
175
+ while len(out) < n:
176
+ out.append(out[-1] if out else "?")
177
+ return out[:n]
178
+
179
+
180
+ def parse_answers(text, n, fallback=None):
181
+ """Raw model output -> exactly n non-empty answers."""
182
+ if not text:
183
+ return list(fallback[:n]) if fallback else ["?"] * n
184
+
185
+ # Prefer the explicit final block the prompt asks for.
186
+ m = None
187
+ for m2 in re.finditer(r"(?:^|\n)\s*(?:final\s+)?answers?\s*:\s*\n?", text, re.I):
188
+ m = m2
189
+ body = text[m.end():] if m else text
190
+
191
+ numbered, raw = [], []
192
+ for ln in body.splitlines():
193
+ if _FENCE.match(ln):
194
+ continue
195
+ mm = re.match(r"^\s*\(?(\d{1,3})\)?[.):\]]\s*(.+)$", ln.strip())
196
+ if mm:
197
+ val = clean_line(mm.group(2))
198
+ if val and not _CHATTY.match(val):
199
+ numbered.append((int(mm.group(1)), val))
200
+ c = clean_line(ln)
201
+ if c and not _CHATTY.match(c):
202
+ raw.append(c)
203
+
204
+ # If the model numbered its answers, trust those labels for placement.
205
+ if len(numbered) >= n:
206
+ by_label = {}
207
+ for lab, val in numbered:
208
+ by_label[lab] = val # last write wins (models restate)
209
+ labs = sorted(by_label)
210
+ if len(labs) >= n:
211
+ return [by_label[l] for l in labs[:n]]
212
+
213
+ return fit_to_n(raw, n, fallback)
214
+
215
+
216
+ def fit_to_n(items, n, fallback=None):
217
+ items = [i for i in items if i and i.strip()]
218
+ if len(items) > n:
219
+ # Take the LAST n. The prompt asks for reasoning first and the answers
220
+ # last, so when there is no ANSWERS: marker to slice on, the tail is the
221
+ # answer block and the head is reasoning prose.
222
+ items = items[-n:]
223
+ while len(items) < n:
224
+ if fallback and len(items) < len(fallback):
225
+ items.append(fallback[len(items)])
226
+ else:
227
+ items.append(items[-1] if items else "?")
228
+ return items[:n]
229
+
230
+
231
+ def norm(s):
232
+ s = unicodedata.normalize("NFC", (s or "").strip().lower())
233
+ s = re.sub(r"\s+", " ", s)
234
+ return s.strip(" .!?;:,")
235
+
236
+
237
+ # ===========================================================================
238
+ # chrF (inline, dependency-free) -- used only to pick the most "central"
239
+ # candidate when self-consistency voting has no majority. sacrebleu is not
240
+ # guaranteed to be importable inside the sandbox.
241
+ # ===========================================================================
242
+
243
+ def _ngrams(s, k):
244
+ s = re.sub(r"\s+", "", s)
245
+ return Counter(s[i:i + k] for i in range(len(s) - k + 1)) if len(s) >= k else Counter()
246
+
247
+
248
+ def chrf_sim(hyp, ref, order=6, beta=2.0):
249
+ if not hyp or not ref:
250
+ return 0.0
251
+ ps, rs = [], []
252
+ for k in range(1, order + 1):
253
+ h, r = _ngrams(hyp, k), _ngrams(ref, k)
254
+ if not h or not r:
255
+ continue
256
+ overlap = sum((h & r).values())
257
+ ps.append(overlap / max(1, sum(h.values())))
258
+ rs.append(overlap / max(1, sum(r.values())))
259
+ if not ps:
260
+ return 0.0
261
+ p, r = sum(ps) / len(ps), sum(rs) / len(rs)
262
+ if p + r == 0:
263
+ return 0.0
264
+ b2 = beta * beta
265
+ return (1 + b2) * p * r / (b2 * p + r)
266
+
267
+
268
+ def vote(cands, anchor=None):
269
+ """Pick one answer for an item, given the greedy answer plus samples.
270
+
271
+ `anchor` is the greedy (temperature-0) answer and is the default. Sampled
272
+ answers may only displace it when at least two of them agree on the same
273
+ normalised form AND that form has strictly more support than the anchor's.
274
+
275
+ This asymmetry is empirically necessary, not decorative. An earlier version
276
+ treated all candidates equally and fell back to "most central by chrF" when
277
+ no majority existed. With only a handful of samples that fallback is
278
+ ill-defined -- with two candidates the pairwise chrF is symmetric, so it
279
+ degenerated to picking the shorter string -- and it replaced the greedy
280
+ answer with a temperature-0.7 sample about half the time. Measured on the
281
+ mock set that cost 4x exact match (EM 0.044 -> 0.011). Anchoring makes
282
+ voting monotone: it can only fire on genuine agreement.
283
+ """
284
+ cands = [c for c in cands if c and c.strip()]
285
+ if anchor is None:
286
+ anchor = cands[0] if cands else "?"
287
+ if len(cands) < 3:
288
+ return anchor
289
+
290
+ groups = defaultdict(list)
291
+ for c in cands:
292
+ groups[norm(c)].append(c)
293
+
294
+ anchor_support = len(groups.get(norm(anchor), []))
295
+ best_key, best_n = None, 0
296
+ for k, v in groups.items():
297
+ if len(v) > best_n:
298
+ best_key, best_n = k, len(v)
299
+
300
+ if best_key is not None and best_n >= 2 and best_n > anchor_support:
301
+ return Counter(groups[best_key]).most_common(1)[0][0]
302
+ return anchor
303
+
304
+
305
+ _OPT_LINE = re.compile(r"^[ \t]*([A-Za-z])[.)]\s+(.+)$", re.M)
306
+ _ITEM_LINE = re.compile(r"^[ \t]*(\d{1,3})[.)]\s+(.+)$", re.M)
307
+
308
+
309
+ def parse_matching_block(context):
310
+ """For match_letters: the numbered items and the lettered options."""
311
+ items = [(int(a), b.strip()) for a, b in _ITEM_LINE.findall(context or "")]
312
+ opts = [(a, b.strip()) for a, b in _OPT_LINE.findall(context or "")]
313
+ seen = set()
314
+ items = [x for x in items if not (x[0] in seen or seen.add(x[0]))]
315
+ seen = set()
316
+ opts = [x for x in opts if not (x[0] in seen or seen.add(x[0]))]
317
+ return items, opts
318
+
319
+
320
+ def best_assignment(score):
321
+ """Max-weight one-to-one assignment. scipy if present, else greedy+swaps."""
322
+ n, m = len(score), len(score[0])
323
+ try:
324
+ from scipy.optimize import linear_sum_assignment
325
+ import numpy as _np
326
+ r, c = linear_sum_assignment(-_np.array(score))
327
+ return list(c)
328
+ except Exception:
329
+ pass
330
+ used, out = set(), [0] * n
331
+ order = sorted(range(n), key=lambda i: -(max(score[i]) - sorted(score[i])[-2]
332
+ if m > 1 else 0))
333
+ for i in order:
334
+ j = max((j for j in range(m) if j not in used),
335
+ key=lambda j: score[i][j], default=0)
336
+ used.add(j)
337
+ out[i] = j
338
+ for _ in range(4): # local 2-swaps
339
+ improved = False
340
+ for a in range(n):
341
+ for b in range(a + 1, n):
342
+ cur = score[a][out[a]] + score[b][out[b]]
343
+ alt = score[a][out[b]] + score[b][out[a]]
344
+ if alt > cur + 1e-9:
345
+ out[a], out[b] = out[b], out[a]
346
+ improved = True
347
+ if not improved:
348
+ break
349
+ return out
350
+
351
+
352
+ def repair_bijection(answers):
353
+ """match_letters answers are usually a permutation of the option letters.
354
+
355
+ When every answer is a single letter and there are as many items as
356
+ distinct letters available, duplicates are certainly wrong. Reassign the
357
+ duplicated slots to the unused letters. Strictly guarded so it is a no-op
358
+ on anything that isn't this shape.
359
+ """
360
+ if len(answers) < 3:
361
+ return answers
362
+ if not all(re.fullmatch(r"[A-Z]", a or "") for a in answers):
363
+ return answers
364
+ n = len(answers)
365
+ universe = [chr(ord("A") + i) for i in range(n)]
366
+ if len(set(answers)) == n:
367
+ return answers
368
+ unused = [l for l in universe if l not in set(answers)]
369
+ if not unused:
370
+ return answers
371
+ seen, out = set(), []
372
+ for a in answers:
373
+ if a in seen and unused:
374
+ out.append(unused.pop(0))
375
+ else:
376
+ seen.add(a)
377
+ out.append(a)
378
+ return out
379
+
380
+
381
+ # ===========================================================================
382
+ # Prompting
383
+ # ===========================================================================
384
+
385
+ SYSTEM = (
386
+ "You are a gold medallist at the International Linguistics Olympiad.\n"
387
+ "Each problem gives data from a language you have never seen. Everything "
388
+ "you need is in the problem itself; no outside knowledge is required or "
389
+ "allowed.\n"
390
+ "Method: line up the given examples, segment the words, identify the "
391
+ "recurring morphemes and the rules that order them, check your rules "
392
+ "against EVERY example, then apply them to the items asked for.\n"
393
+ "Be concise while reasoning. Then output a final block that begins with a "
394
+ "line containing exactly ANSWERS: followed by one answer per line, in the "
395
+ "order asked, with no numbering, no commentary and no blank lines.\n"
396
+ "Give your best guess for every item. Never leave one blank."
397
+ )
398
+
399
+
400
+ # Exact match is half the score, so the answer's *form* matters as much as its
401
+ # content. test.csv states the task type, so say precisely what a well-formed
402
+ # answer looks like. Unknown/absent types simply get no hint.
403
+ TASK_HINTS = {
404
+ "translation": "Each answer is the translation alone -- no source text, no "
405
+ "gloss, no notes, no quotation marks.",
406
+ "match_letters": "Each answer is a single capital letter identifying the "
407
+ "match for that numbered item. Every letter is used "
408
+ "exactly once, so no letter may repeat.",
409
+ "fill_blanks": "Each answer is only the missing form that belongs in that "
410
+ "blank -- not the whole line, not the gloss.",
411
+ "text_to_num": "Each answer is written in digits only (e.g. 111).",
412
+ "num_to_text": "Each answer is the number written out in the problem "
413
+ "language, words only.",
414
+ }
415
+
416
+
417
+ def build_prompt(row, n):
418
+ hint = TASK_HINTS.get((row.get("task_type") or "").strip().lower(), "")
419
+ return (
420
+ f"{row['context'].strip()}\n\n{row['query'].strip()}\n\n"
421
+ f"There are exactly {n} item{'s' if n != 1 else ''} to answer."
422
+ + (f" {hint}" if hint else "") +
423
+ f"\nAfter your reasoning, write ANSWERS: on its own line and then exactly "
424
+ f"{n} line{'s' if n != 1 else ''}, one answer per item, in order."
425
+ )
426
+
427
+
428
+ EXPLAIN_SYSTEM = (
429
+ "You explain International Linguistics Olympiad solutions to a human judge. "
430
+ "Given a problem and the answers produced, state the key rules of the "
431
+ "language that justify them: the relevant morphemes, word order and any "
432
+ "sound changes. Be specific and concise (2-4 sentences or a few short "
433
+ "bullets). Do not restate the reasoning as a stream of thought."
434
+ )
435
+
436
+
437
+ def build_explain_prompt(row, answers):
438
+ return (
439
+ f"{row['context'].strip()}\n\n{row['query'].strip()}\n\n"
440
+ f"Answers given:\n" + "\n".join(f"- {a}" for a in answers) +
441
+ "\n\nBriefly explain the linguistic rules behind these answers."
442
+ )
443
+
444
+
445
+ # ===========================================================================
446
+ # Main
447
+ # ===========================================================================
448
+
449
+ def dev_score(preds):
450
+ """Offline diagnostic: score against a gold file when IOL_GOLD is set.
451
+
452
+ Never runs on the platform (the answers are hidden, so the variable is
453
+ unset there); it exists so one benchmark run reveals the whole learning
454
+ curve -- greedy, then after each self-consistency pass -- instead of a
455
+ single final number.
456
+ """
457
+ gold_path = os.environ.get("IOL_GOLD")
458
+ if not gold_path or not os.path.exists(gold_path):
459
+ return
460
+ try:
461
+ import ast
462
+
463
+ import pandas as pd
464
+ g = pd.read_csv(gold_path, dtype=str)
465
+ ems, cfs = [], []
466
+ for _, r in g.iterrows():
467
+ gold = ast.literal_eval(r["answer"])
468
+ p = preds.get(str(r["id"]), [])
469
+ p = list(p)[:len(gold)] + [""] * max(0, len(gold) - len(p))
470
+ for gi, pi in zip(gold, p):
471
+ alts = gi if isinstance(gi, (list, tuple)) else [gi]
472
+ alts = [str(a) for a in alts]
473
+ ems.append(1.0 if any(pi.strip() == a.strip() for a in alts) else 0.0)
474
+ cfs.append(max(chrf_sim(pi, a) for a in alts))
475
+ em = sum(ems) / max(1, len(ems))
476
+ cf = sum(cfs) / max(1, len(cfs))
477
+ log(f" [dev] EM={em:.4f} chrF~={cf:.4f} score~={(em * cf) ** 0.5:.4f} "
478
+ f"over {len(ems)} items")
479
+ except Exception as e:
480
+ log(f" [dev] scoring failed: {type(e).__name__}: {e}")
481
+
482
+
483
+ def write_submission(path, ids, preds, explanations=None):
484
+ import pandas as pd
485
+ rows = []
486
+ for i in ids:
487
+ rec = {"id": i, "pred": json.dumps(preds[i], ensure_ascii=False)}
488
+ if explanations is not None:
489
+ rec["explanation"] = explanations.get(i, "")
490
+ rows.append(rec)
491
+ pd.DataFrame(rows).to_csv(path, index=False)
492
+
493
+
494
+ def main():
495
+ import pandas as pd
496
+
497
+ df = pd.read_csv(TEST_CSV, dtype=str).fillna("")
498
+ ids = [str(x) for x in df["id"].tolist()]
499
+ ns = [detect_n_items(r.get("query", ""), r.get("task_type", ""), r.get("context", ""))
500
+ for _, r in df.iterrows()]
501
+ total_items = sum(ns)
502
+ log(f"loaded {len(df)} problems, {total_items} items "
503
+ f"(min={min(ns)} max={max(ns)} mean={total_items / len(ns):.1f})")
504
+
505
+ srcs = {i: extract_item_sources(r.get("query", ""), n)
506
+ for i, (_, r), n in zip(ids, df.iterrows(), ns)}
507
+
508
+ # --- 1. Baseline submission on disk before anything can go wrong --------
509
+ preds = {i: list(srcs[i]) for i in ids}
510
+ explanations = {i: "" for i in ids} if WANT_EXPLANATION else None
511
+ write_submission(OUT_CSV, ids, preds, explanations)
512
+ log(f"wrote placeholder {OUT_CSV} ({len(ids)} rows)")
513
+
514
+ # --- 2. Load model -----------------------------------------------------
515
+ import torch
516
+ from transformers import (AutoTokenizer, AutoModelForCausalLM,
517
+ StoppingCriteria, StoppingCriteriaList)
518
+
519
+ class Deadline(StoppingCriteria):
520
+ """Abort generation on wall-clock, checked every token.
521
+
522
+ Without this the budget is only checked between batches, so a batch
523
+ started near the limit runs past it and the platform kills the process.
524
+ """
525
+
526
+ def __init__(self, stop_at):
527
+ self.stop_at = stop_at
528
+
529
+ def __call__(self, input_ids, scores, **kw):
530
+ return time.time() > self.stop_at
531
+
532
+ log("loading tokenizer/model ...")
533
+ tok = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
534
+ if tok.pad_token is None:
535
+ tok.pad_token = tok.eos_token
536
+ tok.padding_side = "left"
537
+
538
+ # Pin every layer to the GPU. device_map="auto" is free to spill layers to
539
+ # CPU when it thinks VRAM is tight, and a couple of offloaded layers make
540
+ # generation ~100x slower without any error -- the worst kind of failure
541
+ # here. Falling back to "auto" only if the explicit placement fails.
542
+ def _load(dev_map):
543
+ # transformers 4.44 (the sandbox) wants torch_dtype=; 5.x renamed it to
544
+ # dtype=. Accept either so the same file runs in both.
545
+ try:
546
+ return AutoModelForCausalLM.from_pretrained(
547
+ MODEL_ID, torch_dtype=torch.float16, device_map=dev_map,
548
+ trust_remote_code=True).eval()
549
+ except TypeError:
550
+ return AutoModelForCausalLM.from_pretrained(
551
+ MODEL_ID, dtype=torch.float16, device_map=dev_map,
552
+ trust_remote_code=True).eval()
553
+
554
+ try:
555
+ model = _load({"": 0} if torch.cuda.is_available() else "auto")
556
+ except Exception as e:
557
+ log(f"pinned load failed ({type(e).__name__}: {e}); falling back to auto")
558
+ model = _load("auto")
559
+
560
+ devs = set(str(p.device) for p in model.parameters())
561
+ log(f"model ready on {sorted(devs)} ({left():.0f}s of budget left)")
562
+ if any(d.startswith("cpu") or d == "meta" for d in devs):
563
+ log("WARNING: part of the model is off-GPU; generation will be very slow")
564
+ if torch.cuda.is_available():
565
+ log(f" VRAM allocated {torch.cuda.memory_allocated()/1e9:.2f} GB / "
566
+ f"{torch.cuda.get_device_properties(0).total_memory/1e9:.1f} GB")
567
+
568
+ prompts = []
569
+ for (_, r), n in zip(df.iterrows(), ns):
570
+ if BASELINE_MODE:
571
+ msgs = [{"role": "system", "content":
572
+ "You solve International Linguistics Olympiad problems. "
573
+ "Answer every numbered item. Put each answer on its own line, "
574
+ "in order, with no numbering and no extra text."},
575
+ {"role": "user", "content":
576
+ f"{r['context'].strip()}\n\n{r['query'].strip()}"}]
577
+ else:
578
+ msgs = [{"role": "system", "content": SYSTEM},
579
+ {"role": "user", "content": build_prompt(r, n)}]
580
+ prompts.append(tok.apply_chat_template(msgs, tokenize=False,
581
+ add_generation_prompt=True))
582
+
583
+ batch_size = 1 if BASELINE_MODE else int(os.environ.get("IOL_BATCH", "4"))
584
+
585
+ def generate(texts, max_new, sample, temp=0.7):
586
+ """Batched generation with OOM backoff. Returns list of strings."""
587
+ nonlocal batch_size
588
+ out = [""] * len(texts)
589
+ order = sorted(range(len(texts)), key=lambda i: len(texts[i]))
590
+ i = 0
591
+ while i < len(order):
592
+ if left() < 25:
593
+ log(" out of time inside generate(); returning partial")
594
+ break
595
+ idx = order[i:i + batch_size]
596
+ chunk = [texts[j] for j in idx]
597
+ try:
598
+ enc = tok(chunk, return_tensors="pt", padding=True,
599
+ truncation=True, max_length=6144).to(model.device)
600
+ # repetition_penalty=1.0 EXPLICITLY. Qwen2.5-14B-Instruct-AWQ
601
+ # ships generation_config.json with repetition_penalty=1.05,
602
+ # and unlike temperature/top_p/top_k (which greedy ignores, and
603
+ # which transformers warns about) a repetition penalty IS
604
+ # applied under greedy decoding -- silently, with no warning.
605
+ # 34% of the public gold answers repeat a letter 3+ times
606
+ # (agglutinative morphology like 'ɨmpʼuhurʼu'), so a 5% penalty
607
+ # pushes the model off exactly the strings we need.
608
+ kw = dict(max_new_tokens=max_new, pad_token_id=tok.pad_token_id,
609
+ repetition_penalty=1.0,
610
+ stopping_criteria=StoppingCriteriaList(
611
+ [Deadline(DEADLINE - 10)]))
612
+ if sample:
613
+ kw.update(do_sample=True, temperature=temp, top_p=0.95)
614
+ else:
615
+ kw.update(do_sample=False)
616
+ with torch.no_grad():
617
+ o = model.generate(**enc, **kw)
618
+ for k, j in enumerate(idx):
619
+ out[j] = tok.decode(o[k][enc["input_ids"].shape[1]:],
620
+ skip_special_tokens=True)
621
+ i += batch_size
622
+ except torch.cuda.OutOfMemoryError:
623
+ torch.cuda.empty_cache()
624
+ if batch_size == 1:
625
+ log(" OOM at batch=1; skipping this item")
626
+ i += 1
627
+ else:
628
+ batch_size = max(1, batch_size // 2)
629
+ log(f" OOM -> batch_size={batch_size}")
630
+ except Exception as e: # never die mid-run
631
+ log(f" generate error: {type(e).__name__}: {e}")
632
+ i += batch_size
633
+ return out
634
+
635
+ def solve_matching(row, n):
636
+ """Score every (item, option) pair and take the best one-to-one assignment.
637
+
638
+ Free-form generation fails badly here: measured on the benchmark the
639
+ model just emits the option labels in order (A, B, C, ... == the
640
+ identity permutation), which is a *valid* permutation so no repair
641
+ fires, and it scores ~0. Asking for one letter at a time and reading
642
+ the next-token distribution turns the task into an assignment problem
643
+ the model is actually good at, and the one-to-one constraint is then
644
+ enforced exactly rather than hoped for.
645
+ """
646
+ items, opts = parse_matching_block(row.get("context", ""))
647
+ if len(items) < 3 or len(opts) < 3 or len(items) != n:
648
+ return None
649
+ letters = [o[0] for o in opts]
650
+ # token id for each option letter, bare and space-prefixed
651
+ cand_ids = []
652
+ for L in letters:
653
+ ids = set()
654
+ for form in (L, " " + L):
655
+ t = tok.encode(form, add_special_tokens=False)
656
+ if t:
657
+ ids.add(t[0])
658
+ cand_ids.append(sorted(ids))
659
+
660
+ ctx = row["context"].strip()
661
+ prompts_m = []
662
+ for num, itext in items:
663
+ msgs = [
664
+ {"role": "system", "content":
665
+ "You match items to their correct counterparts in a "
666
+ "linguistics problem. Reply with one option letter only."},
667
+ {"role": "user", "content":
668
+ f"{ctx}\n\nWhich lettered option corresponds to item {num} "
669
+ f"({itext})? Reply with the option letter only."},
670
+ ]
671
+ prompts_m.append(tok.apply_chat_template(
672
+ msgs, tokenize=False, add_generation_prompt=True))
673
+
674
+ score = []
675
+ bs = 4
676
+ for s0 in range(0, len(prompts_m), bs):
677
+ if left() < 30:
678
+ return None
679
+ chunk = prompts_m[s0:s0 + bs]
680
+ enc = tok(chunk, return_tensors="pt", padding=True,
681
+ truncation=True, max_length=6144).to(model.device)
682
+ with torch.no_grad():
683
+ logits = model(**enc).logits[:, -1, :].float()
684
+ logprobs = torch.log_softmax(logits, dim=-1)
685
+ for b in range(len(chunk)):
686
+ score.append([max(logprobs[b, i].item() for i in ids)
687
+ for ids in cand_ids])
688
+ col = best_assignment(score)
689
+ return [letters[c] for c in col]
690
+
691
+ # --- 3. Pass 1: greedy, guarantees a full answer set --------------------
692
+ # Size the reasoning budget to the actual problem count. Measured on the
693
+ # eval hardware (T4, 14B AWQ, batch 4) throughput is ~32 tok/s, so the whole
694
+ # 30 minutes buys only ~50k generated tokens. With ~16 problem blocks that
695
+ # affords full-length reasoning; if the platform instead ships one row per
696
+ # sub-question (~90 rows) a fixed 900-token budget would not even finish a
697
+ # single pass. Spend at most ~40% of what's left on pass 1.
698
+ TOK_PER_S = float(os.environ.get("IOL_TOKS", "30"))
699
+ adaptive = int(0.40 * max(1.0, left()) * TOK_PER_S / max(1, len(df)))
700
+ max_new = max(192, min(MAX_NEW, adaptive))
701
+ log(f"reasoning budget: {max_new} new tokens/problem "
702
+ f"(adaptive={adaptive}, cap={MAX_NEW}, {len(df)} problems)")
703
+
704
+ t = time.time()
705
+ texts = generate(prompts, max_new=max_new, sample=False)
706
+ pass1_cost = time.time() - t
707
+ samples = {i: [] for i in ids}
708
+ n_matched = 0
709
+ solved_ids = set() # rows answered by the assignment solver, vote-exempt
710
+ for (i, n, txt), (_, row) in zip(zip(ids, ns, texts), df.iterrows()):
711
+ if BASELINE_MODE:
712
+ # literally the organizers' parse: every non-empty stripped line,
713
+ # however many there are. No cleaning, no fallback, no forcing.
714
+ preds[i] = [ln.strip() for ln in (txt or "").splitlines() if ln.strip()]
715
+ # v12: THE single change on top of the 0.2245 v8 config. Free-form
716
+ # generation on match_letters emits the identity permutation
717
+ # (A, B, C, ...) and scores ~0 (measured: 7 parser variants all
718
+ # gave EM 0.0000 on these, exonerating the parser). Replace only
719
+ # this task type with the assignment solver; every other row is
720
+ # byte-identical to v8.
721
+ if (row.get("task_type") or "").strip().lower() == "match_letters":
722
+ try:
723
+ mm_ = solve_matching(row, n)
724
+ if mm_ and len(mm_) == n:
725
+ preds[i] = mm_
726
+ n_matched += 1
727
+ solved_ids.add(i)
728
+ except Exception as e:
729
+ log(f" matching solver failed on {i}: {type(e).__name__}: {e}")
730
+ samples[i].append(preds[i])
731
+ continue
732
+ a = repair_bijection(parse_answers(txt, n, srcs[i]))
733
+ # match_letters: free-form generation emits the identity permutation
734
+ # (A, B, C, ...) and scores ~0, so solve it as an assignment instead.
735
+ if (row.get("task_type") or "").strip().lower() == "match_letters":
736
+ try:
737
+ mm_ = solve_matching(row, n)
738
+ if mm_ and len(mm_) == n:
739
+ a = mm_
740
+ n_matched += 1
741
+ except Exception as e:
742
+ log(f" matching solver failed on {i}: {type(e).__name__}: {e}")
743
+ preds[i] = a
744
+ samples[i].append(a)
745
+ if n_matched:
746
+ log(f"assignment solver used on {n_matched} match_letters problem(s)")
747
+ write_submission(OUT_CSV, ids, preds, explanations)
748
+ # How often did reasoning run past the token budget before the model got to
749
+ # its ANSWERS: block? Those problems fall back to salvaged lines, so a high
750
+ # count means max_new is too small rather than the model being wrong.
751
+ no_block = sum(1 for txt in texts
752
+ if not re.search(r"answers?\s*:", txt or "", re.I))
753
+ empty = sum(1 for txt in texts if not (txt or "").strip())
754
+ log(f"pass 1 (greedy) done in {pass1_cost:.0f}s -> submission written "
755
+ f"({no_block}/{len(texts)} without an ANSWERS: block, {empty} empty)")
756
+ dev_score(preds)
757
+
758
+ # --- 4. Self-consistency passes while budget allows ---------------------
759
+ reserve = 0.0
760
+ if WANT_EXPLANATION:
761
+ reserve = min(300.0, 0.25 * pass1_cost + 60) # explanations are short
762
+ n_extra = 0
763
+ while left() - reserve > pass1_cost * 1.25 and n_extra < MAX_SAMPLES:
764
+ n_extra += 1
765
+ log(f"self-consistency pass {n_extra} ({left():.0f}s left)")
766
+ texts = generate(prompts, max_new=max_new, sample=True, temp=SAMPLE_TEMP)
767
+ for i, n, txt in zip(ids, ns, texts):
768
+ if txt:
769
+ samples[i].append(repair_bijection(parse_answers(txt, n, srcs[i])))
770
+ for i, n in zip(ids, ns):
771
+ if i in solved_ids:
772
+ # the sampled passes re-emit the identity permutation; voting
773
+ # over them would undo the solver exactly where it matters
774
+ continue
775
+ # samples[i][0] is the greedy pass; it anchors every item.
776
+ if len(samples[i]) >= 3:
777
+ greedy = samples[i][0]
778
+ preds[i] = repair_bijection(
779
+ [vote([s[k] for s in samples[i] if k < len(s)],
780
+ anchor=greedy[k] if k < len(greedy) else None)
781
+ for k in range(n)])
782
+ write_submission(OUT_CSV, ids, preds, explanations)
783
+ log(f" voted over {n_extra + 1} samples (greedy-anchored) -> written")
784
+ dev_score(preds)
785
+
786
+ # --- 5. Explanations for the jury track ---------------------------------
787
+ if WANT_EXPLANATION and left() > 60:
788
+ log(f"generating explanations ({left():.0f}s left)")
789
+ ex_prompts = []
790
+ for (_, r), i in zip(df.iterrows(), ids):
791
+ msgs = [{"role": "system", "content": EXPLAIN_SYSTEM},
792
+ {"role": "user", "content": build_explain_prompt(r, preds[i])}]
793
+ ex_prompts.append(tok.apply_chat_template(
794
+ msgs, tokenize=False, add_generation_prompt=True))
795
+ ex = generate(ex_prompts, max_new=200, sample=False)
796
+ for i, e in zip(ids, ex):
797
+ e = re.sub(r"\s+", " ", (e or "").strip())
798
+ if e:
799
+ explanations[i] = e[:1200]
800
+ write_submission(OUT_CSV, ids, preds, explanations)
801
+ log("explanations written")
802
+
803
+ # --- 6. Final integrity check ------------------------------------------
804
+ bad = [i for i, n in zip(ids, ns) if len(preds[i]) != n or any(
805
+ not str(x).strip() for x in preds[i])]
806
+ if bad:
807
+ log(f"repairing {len(bad)} malformed rows")
808
+ for i, n in zip(ids, ns):
809
+ preds[i] = fit_to_n([x for x in preds[i] if str(x).strip()], n, srcs[i])
810
+ write_submission(OUT_CSV, ids, preds, explanations)
811
+
812
+ log(f"DONE. {len(ids)} rows, {sum(len(v) for v in preds.values())} answers, "
813
+ f"{time.time() - T0:.0f}s elapsed")
814
+
815
+
816
+ if __name__ == "__main__":
817
+ main()