609 lines
27 KiB
Python
609 lines
27 KiB
Python
# script.py — Qwen3-14B-AWQ + OUR decomposition pipeline (Srikar's offline
|
|
# wheelhouse used ONLY as the loading mechanism, not his prompt).
|
|
# =============================================================================
|
|
# WHAT CHANGED vs the 0.104 baseline (single conceptual variable = the model):
|
|
# 1. Base model: Qwen2.5-14B-bnb-4bit -> Qwen3-14B-AWQ.
|
|
# 2. Dependency install: instead of `pip install --no-deps bitsandbytes`
|
|
# (base env), we install the Qwen3 stack from BUNDLED wheels with
|
|
# `pip install --no-index --no-deps --target <RUNTIME_DIR>` and prepend
|
|
# that dir to sys.path. This NEVER touches the network and NEVER mutates
|
|
# base site-packages, so the grader's later hf_hub_download (metric.py)
|
|
# runs on the pristine base huggingface_hub -- the RemoteDisconnected
|
|
# class of failure cannot recur.
|
|
# 3. Chat template: pass enable_thinking=False (Qwen3 supports it; harmless
|
|
# on models that ignore it). We keep OUR own decomposition reasoning in
|
|
# the prompt rather than paying for Qwen3's <think> phase.
|
|
# Everything else -- symbolic evidence, FINAL ANSWERS contract, safe
|
|
# arithmetic, explanations, per-row crash safety, dynamic token budget,
|
|
# guaranteed one row per id -- is IDENTICAL to the proven 0.104 pipeline.
|
|
# match_letters bijection decoding is deliberately NOT added here; that is the
|
|
# next, separate experiment.
|
|
# =============================================================================
|
|
import os
|
|
import atexit
|
|
from pathlib import Path
|
|
_ORIGINAL_HF_HUB_OFFLINE = os.environ.get("HF_HUB_OFFLINE")
|
|
_ORIGINAL_TRANSFORMERS_OFFLINE = os.environ.get("TRANSFORMERS_OFFLINE")
|
|
def _restore_offline_env_vars():
|
|
for key, original in (("HF_HUB_OFFLINE", _ORIGINAL_HF_HUB_OFFLINE),
|
|
("TRANSFORMERS_OFFLINE", _ORIGINAL_TRANSFORMERS_OFFLINE)):
|
|
if original is None:
|
|
os.environ.pop(key, None)
|
|
else:
|
|
os.environ[key] = original
|
|
atexit.register(_restore_offline_env_vars)
|
|
os.environ.setdefault("HF_HUB_OFFLINE", "1")
|
|
os.environ.setdefault("TRANSFORMERS_OFFLINE", "1")
|
|
import subprocess, sys
|
|
import importlib
|
|
import importlib.metadata
|
|
SCRIPT_DIR = Path(__file__).resolve().parent
|
|
def emergency_submission_csv(reason, rows_so_far=None):
|
|
try:
|
|
import pandas as pd
|
|
if rows_so_far:
|
|
pd.DataFrame(rows_so_far).to_csv("submission.csv", index=False)
|
|
return
|
|
try:
|
|
df = pd.read_csv("/tmp/data/test.csv", dtype=str).fillna("")
|
|
ids = df["id"].tolist()
|
|
except Exception:
|
|
ids = []
|
|
import json as _json
|
|
rows = [{"id": i, "pred": _json.dumps([""]),
|
|
"explanation": f"EMERGENCY FALLBACK: {str(reason)[:150]}"} for i in ids]
|
|
pd.DataFrame(rows, columns=["id", "pred", "explanation"]).to_csv("submission.csv", index=False)
|
|
except Exception:
|
|
try:
|
|
with open("submission.csv", "w") as f:
|
|
f.write("id,pred,explanation\n")
|
|
except Exception:
|
|
pass
|
|
def write_submission_csv(rows_list):
|
|
import csv as _csv
|
|
tmp_path = "submission.csv.tmp"
|
|
with open(tmp_path, "w", newline="", encoding="utf-8") as f:
|
|
w = _csv.DictWriter(f, fieldnames=["id", "pred", "explanation"])
|
|
w.writeheader()
|
|
for row in rows_list:
|
|
w.writerow(row)
|
|
os.replace(tmp_path, "submission.csv")
|
|
# =============================================================================
|
|
# OFFLINE WHEELHOUSE INSTALL (adapted from the public 0.147 submissions).
|
|
# Installs the Qwen3-compatible stack from wheels bundled inside this repo,
|
|
# to an isolated --target dir that we prepend to sys.path. --no-index means
|
|
# pip never contacts the network (the sandbox has no working index anyway);
|
|
# --target means base site-packages is untouched, so scoring stays safe.
|
|
# =============================================================================
|
|
WHEELHOUSE = Path(os.environ.get("QWEN3_WHEELHOUSE", str(SCRIPT_DIR / "wheelhouse")))
|
|
RUNTIME_DIR = Path(os.environ.get("QWEN3_RUNTIME_DIR", "/tmp/qwen3deps"))
|
|
RUNTIME_PACKAGES = {
|
|
"transformers": "4.51.3",
|
|
"tokenizers": "0.21.1",
|
|
"huggingface_hub": "0.30.2",
|
|
"autoawq": "0.2.9",
|
|
}
|
|
RUNTIME_WHEELS = (
|
|
"transformers-4.51.3-py3-none-any.whl",
|
|
"tokenizers-0.21.1-cp39-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl",
|
|
"huggingface_hub-0.30.2-py3-none-any.whl",
|
|
"autoawq-0.2.9-py3-none-any.whl",
|
|
)
|
|
def ensure_runtime_dependencies():
|
|
wheel_paths = [WHEELHOUSE / name for name in RUNTIME_WHEELS]
|
|
missing = [str(p) for p in wheel_paths if not p.is_file()]
|
|
if missing:
|
|
raise FileNotFoundError(f"Missing offline runtime wheels: {missing}")
|
|
marker = RUNTIME_DIR / ".iol-qwen3-runtime-v1"
|
|
if not marker.is_file():
|
|
RUNTIME_DIR.mkdir(parents=True, exist_ok=True)
|
|
subprocess.run(
|
|
[sys.executable, "-m", "pip", "install",
|
|
"--disable-pip-version-check", "--no-index", "--no-deps",
|
|
"--upgrade", "--target", str(RUNTIME_DIR),
|
|
*(str(p) for p in wheel_paths)],
|
|
check=True, timeout=300,
|
|
)
|
|
marker.write_text("offline Qwen3 runtime installed\n", encoding="utf-8")
|
|
runtime_path = str(RUNTIME_DIR)
|
|
if runtime_path in sys.path:
|
|
sys.path.remove(runtime_path)
|
|
sys.path.insert(0, runtime_path)
|
|
importlib.invalidate_caches()
|
|
versions = {}
|
|
for pkg in RUNTIME_PACKAGES:
|
|
try:
|
|
versions[pkg] = importlib.metadata.version(pkg)
|
|
except importlib.metadata.PackageNotFoundError:
|
|
versions[pkg] = "missing"
|
|
print(f"offline runtime active: {versions}", flush=True)
|
|
try:
|
|
ensure_runtime_dependencies()
|
|
except Exception as e:
|
|
emergency_submission_csv(f"wheelhouse install failed: {e}")
|
|
raise
|
|
import re, json, time, ast as pyast
|
|
import pandas as pd
|
|
import torch
|
|
from transformers import AutoTokenizer, AutoModelForCausalLM
|
|
MODEL_ID = "."
|
|
TIME_LIMIT_S = 30 * 60
|
|
SETUP_BUFFER_S = 360
|
|
start_time = time.time()
|
|
try:
|
|
df = pd.read_csv("/tmp/data/test.csv", dtype=str).fillna("")
|
|
placeholder_rows = [{"id": rid, "pred": json.dumps([""]),
|
|
"explanation": "Placeholder written before model load."}
|
|
for rid in df["id"].tolist()]
|
|
write_submission_csv(placeholder_rows)
|
|
print(f"Pre-load checkpoint written for {len(placeholder_rows)} rows.", flush=True)
|
|
# AWQ backend preflight (diagnostic only, never fatal).
|
|
try:
|
|
from awq.modules.linear import gemm as awq_gemm
|
|
print(f"AWQ backends: extension={awq_gemm.awq_ext is not None}, "
|
|
f"triton={getattr(awq_gemm, 'TRITON_AVAILABLE', None)}", flush=True)
|
|
except Exception as exc:
|
|
print(f"AWQ backend preflight warning: {exc}", flush=True)
|
|
try:
|
|
tok = AutoTokenizer.from_pretrained(MODEL_ID, local_files_only=True)
|
|
print("Tokenizer loaded (fast).", flush=True)
|
|
except Exception as e:
|
|
print(f"Fast tokenizer failed ({e}); falling back to use_fast=False.", flush=True)
|
|
tok = AutoTokenizer.from_pretrained(MODEL_ID, use_fast=False, local_files_only=True)
|
|
print("Tokenizer loaded (slow fallback).", flush=True)
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
MODEL_ID, torch_dtype=torch.float16, device_map="auto", local_files_only=True,
|
|
).eval()
|
|
print(f"Model loaded | memory footprint: {round(model.get_memory_footprint()/1e9, 1)} GB | "
|
|
f"quantized: {getattr(model.config, 'quantization_config', None) is not None}", flush=True)
|
|
except Exception as e:
|
|
emergency_submission_csv(f"tokenizer/model load or test.csv read failed: {e}")
|
|
raise
|
|
n_rows = len(df)
|
|
actual_setup_elapsed = time.time() - start_time
|
|
per_row_budget = max(20, (TIME_LIMIT_S - actual_setup_elapsed) / max(n_rows, 1))
|
|
print(f"Setup took {actual_setup_elapsed:.0f}s | per_row_budget={per_row_budget:.0f}s "
|
|
f"for {n_rows} rows", flush=True)
|
|
# ---- Query parsing ----
|
|
def parse_items(query: str):
|
|
item_pat = re.compile(r"(?m)^\s*(\d+)\s*[.\)]\s*(.*)$")
|
|
matches = list(item_pat.finditer(query))
|
|
if matches:
|
|
preamble = query[:matches[0].start()].strip()
|
|
items = []
|
|
for i, m in enumerate(matches):
|
|
end = matches[i + 1].start() if i + 1 < len(matches) else len(query)
|
|
text = re.sub(r"^\s*\d+\s*[.\)]\s*", "", query[m.start():end].strip())
|
|
items.append(text)
|
|
return preamble, items, True
|
|
rng = re.search(r"[\(\[]?\s*(\d+)\s*(?:[-–—:]|to)\s*(\d+)\s*[\)\]]?", query, flags=re.IGNORECASE)
|
|
if rng:
|
|
lo, hi = int(rng.group(1)), int(rng.group(2))
|
|
if 0 < hi - lo < 100:
|
|
items = []
|
|
for k in range(lo, hi + 1):
|
|
line_match = re.search(rf"(?m)^.*\(\s*{k}\s*\).*$", query)
|
|
if line_match:
|
|
clue = re.sub(rf"\(\s*{k}\s*\)", "", line_match.group(0)).strip()
|
|
clue = re.sub(r"\|\s*\|", "|", clue)
|
|
clue = re.sub(r"\s{2,}", " ", clue).strip(" |")
|
|
items.append(clue if clue else f"the numbered item {k} from the examples above")
|
|
else:
|
|
items.append(f"the numbered item {k} from the examples above")
|
|
return query.strip(), items, True
|
|
csv_nums = re.findall(r"(?m)^\s*(\d+)\s*,\s*(\d+(?:\s*,\s*\d+)*)\s*$", query)
|
|
if csv_nums:
|
|
all_nums = re.findall(r"\d+", " ".join(csv_nums[0]))
|
|
return query.strip(), [f"the numbered item {n}" for n in all_nums], True
|
|
return query.strip(), [], False
|
|
TASK_GUIDANCE = {
|
|
"translation": "give the translated form only, in the language asked.",
|
|
"fill_blanks": "give only the missing form for each blank.",
|
|
"match_letters": "give only the option letter (for example A, B, C).",
|
|
"text_to_num": "give the number in digits.",
|
|
"num_to_text": "give the number written out in words, in the language asked.",
|
|
}
|
|
DEFAULT_GUIDANCE = "give exactly what the instruction asks, nothing else."
|
|
from difflib import SequenceMatcher
|
|
from collections import defaultdict
|
|
def extract_forms_from_context(context: str):
|
|
forms = []
|
|
for line in context.splitlines():
|
|
line = line.strip()
|
|
if not line:
|
|
continue
|
|
pipe_count = line.count("|")
|
|
if 0 < pipe_count <= 3:
|
|
first_field = re.sub(r"^\s*\d+\s*[.\)]\s*", "", line.split("|")[0].strip()).strip()
|
|
if first_field:
|
|
forms.append(first_field)
|
|
elif pipe_count == 0:
|
|
for t in line.split():
|
|
t_clean = re.sub(r"^\s*\d+\s*[.\)]\s*", "", t).strip(".,;:")
|
|
if t_clean and len(t_clean) > 1:
|
|
forms.append(t_clean)
|
|
seen, unique_forms = set(), []
|
|
for f in forms:
|
|
if f not in seen:
|
|
seen.add(f)
|
|
unique_forms.append(f)
|
|
return unique_forms
|
|
def extract_explicit_pairs(context: str):
|
|
pairs = []
|
|
for line in context.splitlines():
|
|
line = line.strip()
|
|
if not (0 < line.count("|") <= 3):
|
|
continue
|
|
fields = [re.sub(r"^\s*\d+\s*[.\)]\s*", "", f.strip()).strip() for f in line.split("|")]
|
|
fields = [f for f in fields if f]
|
|
if len(fields) >= 2:
|
|
pairs.append((fields[0], fields[1]))
|
|
return pairs
|
|
def edit_signature(a: str, b: str):
|
|
sm = SequenceMatcher(None, a, b, autojunk=False)
|
|
all_ops = sm.get_opcodes()
|
|
ops = [op for op in all_ops if op[0] != "equal"]
|
|
if not ops or len(ops) > 2:
|
|
return None
|
|
equal_len = sum((i2 - i1) for tag, i1, i2, j1, j2 in all_ops if tag == "equal")
|
|
if equal_len < 2:
|
|
return None
|
|
tag, i1, i2, j1, j2 = ops[0]
|
|
removed, inserted = a[i1:i2], b[j1:j2]
|
|
if i1 == 0:
|
|
pos = "prefix"
|
|
elif i2 == len(a):
|
|
pos = "suffix"
|
|
else:
|
|
pos = "infix"
|
|
return (pos, removed, inserted)
|
|
def find_transformation_families(pairs):
|
|
groups = defaultdict(list)
|
|
for a, b in pairs:
|
|
if not a or not b or a == b:
|
|
continue
|
|
sig = edit_signature(a, b)
|
|
if sig:
|
|
groups[sig].append((a, b))
|
|
families = []
|
|
for sig, grp in groups.items():
|
|
unique_pairs = list(dict.fromkeys(grp))
|
|
if len(unique_pairs) >= 2:
|
|
pos, removed, inserted = sig
|
|
removed_disp = removed if removed else "(nothing)"
|
|
inserted_disp = inserted if inserted else "(nothing)"
|
|
examples = "; ".join(f"{a}->{b}" for a, b in unique_pairs[:4])
|
|
families.append((len(unique_pairs),
|
|
f"{pos} change: '{removed_disp}' -> '{inserted_disp}' (seen in: {examples})"))
|
|
families.sort(key=lambda x: -x[0])
|
|
return [f for _, f in families]
|
|
def detect_reduplication(forms):
|
|
findings = []
|
|
for w in forms:
|
|
n = len(w)
|
|
found = False
|
|
for length in range(2, n // 2 + 1):
|
|
for start in range(0, n - 2 * length + 1):
|
|
chunk = w[start:start + length]
|
|
nxt = w[start + length:start + 2 * length]
|
|
if chunk == nxt:
|
|
findings.append(f"reduplication in '{w}': '{chunk}' repeated")
|
|
found = True
|
|
break
|
|
if found:
|
|
break
|
|
return findings
|
|
def build_symbolic_evidence(context: str) -> str:
|
|
forms = extract_forms_from_context(context)
|
|
pairs = extract_explicit_pairs(context)
|
|
families = find_transformation_families(pairs) if pairs else []
|
|
redup = detect_reduplication(forms) if forms else []
|
|
lines = []
|
|
if families:
|
|
lines.append("Transformation families found (patterns supported by multiple examples):")
|
|
for f in families[:3]:
|
|
lines.append(f"- {f}")
|
|
if redup:
|
|
lines.append("Reduplication detected:")
|
|
for r in redup[:2]:
|
|
lines.append(f"- {r}")
|
|
if not lines:
|
|
return ""
|
|
return ("\n\nSYMBOLIC EVIDENCE (deterministically computed from the examples above; "
|
|
"may be incomplete -- verify against the examples, do not trust blindly):\n"
|
|
+ "\n".join(lines))
|
|
def build_messages(context, query, task_type):
|
|
preamble, items, count_known = parse_items(query)
|
|
guidance = TASK_GUIDANCE.get(task_type, DEFAULT_GUIDANCE)
|
|
symbolic_evidence = build_symbolic_evidence(context)
|
|
system = (
|
|
"You solve puzzles about a language you have never seen. Everything you "
|
|
"need is in the examples below. Use only the examples, not outside "
|
|
"knowledge of any language. You may meet a task type you have never "
|
|
"seen -- read the instruction and examples, and answer in the same "
|
|
"form they use."
|
|
)
|
|
number_note = ""
|
|
if task_type == "text_to_num":
|
|
number_note = (
|
|
"\n\nAlso add one more line after your answers, exactly like this:\n"
|
|
"COMPUTE: expr1 | expr2\n"
|
|
"where each expr is a plain arithmetic expression (digits, +, -, *, "
|
|
"parentheses only) for that item's value, one per answer, matching "
|
|
"the rule you found."
|
|
)
|
|
options_note = ""
|
|
if task_type == "match_letters":
|
|
options = extract_match_letter_options(context)
|
|
if options:
|
|
options_note = (
|
|
f"\n\nThe only valid answers are: {', '.join(options)}. "
|
|
f"Do not use any other letter."
|
|
)
|
|
if count_known:
|
|
n_items = len(items)
|
|
slots = "\n\n".join(f"Question {i+1}: {it}\nAnswer {i+1}:" for i, it in enumerate(items))
|
|
user = (
|
|
f"EXAMPLES:\n{context.strip()}"
|
|
f"{symbolic_evidence}\n\n"
|
|
f"--- The examples end here. The questions begin below. ---\n\n"
|
|
f"For each question: find the rule that explains ALL the examples above "
|
|
f"(not just one). Check it against every example before answering. "
|
|
f"For this task type, {guidance}\n\n"
|
|
f"{preamble}\n\n{slots}\n\n"
|
|
f"After answering all {n_items} questions, finish with exactly one line, "
|
|
f"all {n_items} answers in order separated by ' | ':\n"
|
|
f"FINAL ANSWERS: answer1 | answer2"
|
|
f"{number_note}"
|
|
f"{options_note}"
|
|
)
|
|
else:
|
|
n_items = None
|
|
user = (
|
|
f"EXAMPLES:\n{context.strip()}"
|
|
f"{symbolic_evidence}\n\n"
|
|
f"--- The examples end here. The question begins below. ---\n\n"
|
|
f"Find the rule that explains ALL the examples above (not just one). "
|
|
f"Check it against every example before answering. "
|
|
f"For this task type, {guidance}\n\n"
|
|
f"{preamble}\n\n"
|
|
f"Answer every item asked above, in order, one per answer. Finish "
|
|
f"with exactly one line, all your answers in order separated by ' | ':\n"
|
|
f"FINAL ANSWERS: answer1 | answer2"
|
|
f"{number_note}"
|
|
f"{options_note}"
|
|
)
|
|
return [{"role": "system", "content": system}, {"role": "user", "content": user}], n_items
|
|
def build_repair_messages(query, n_items, bad_text):
|
|
n_desc = f"exactly {n_items}" if n_items is not None else "one per item asked"
|
|
system = "You reformat answers. Output nothing except the requested line."
|
|
user = (
|
|
f"Question:\n{query.strip()}\n\n"
|
|
f"A previous attempt produced:\n{bad_text[:600]}\n\n"
|
|
f"Extract or restate {n_desc} final answers, in order, as ONE line:\n"
|
|
f"FINAL ANSWERS: answer1 | answer2"
|
|
)
|
|
return [{"role": "system", "content": system}, {"role": "user", "content": user}]
|
|
_ALLOWED_BINOPS = (pyast.Add, pyast.Sub, pyast.Mult)
|
|
def safe_arithmetic(expr: str):
|
|
try:
|
|
tree = pyast.parse(expr.strip(), mode="eval")
|
|
except Exception:
|
|
return None
|
|
def _eval(node):
|
|
if isinstance(node, pyast.Expression):
|
|
return _eval(node.body)
|
|
if isinstance(node, pyast.Constant) and isinstance(node.value, (int, float)):
|
|
return node.value
|
|
if isinstance(node, pyast.BinOp) and isinstance(node.op, _ALLOWED_BINOPS):
|
|
left, right = _eval(node.left), _eval(node.right)
|
|
if left is None or right is None:
|
|
return None
|
|
if isinstance(node.op, pyast.Add): return left + right
|
|
if isinstance(node.op, pyast.Sub): return left - right
|
|
if isinstance(node.op, pyast.Mult): return left * right
|
|
if isinstance(node, pyast.UnaryOp) and isinstance(node.op, pyast.USub):
|
|
v = _eval(node.operand)
|
|
return -v if v is not None else None
|
|
return None
|
|
return _eval(tree)
|
|
def clean_answer(a: str) -> str:
|
|
a = re.sub(r"(?i)^\s*(the\s+)?(final\s+)?answer\s*\d*\s*(is)?\s*:\s*", "", a).strip()
|
|
a = re.sub(r"(?i)^\s*is\s*:\s*", "", a).strip()
|
|
a = a.strip("* ")
|
|
return a.strip(" .\"'“”‘’")
|
|
def extract(text):
|
|
m = list(re.finditer(r"final answers?\s*:?\s*\**", text, flags=re.IGNORECASE))
|
|
if m:
|
|
tail = text[m[-1].end():]
|
|
stop = re.search(r"(?i)compute\s*:", tail)
|
|
if stop:
|
|
tail = tail[:stop.start()]
|
|
tail = tail.replace("**", " ").strip()
|
|
candidate = " ".join(tail.splitlines())
|
|
parts = [clean_answer(p) for p in candidate.split("|") if p.strip()]
|
|
if parts:
|
|
return parts, m[-1].start()
|
|
lines = [ln.strip() for ln in text.splitlines() if ln.strip()]
|
|
fallback = []
|
|
for ln in lines:
|
|
ln_clean = re.sub(r"^\s*\d+\s*[.\)]\s*", "", ln)
|
|
if "|" in ln_clean:
|
|
fallback.extend(clean_answer(p) for p in ln_clean.split("|") if p.strip())
|
|
else:
|
|
fallback.append(clean_answer(ln_clean))
|
|
return fallback, None
|
|
def extract_compute_overrides(text, n_answers):
|
|
m = re.search(r"compute\s*:\s*(.+)", text, flags=re.IGNORECASE)
|
|
if not m:
|
|
return {}
|
|
exprs = [e.strip() for e in m.group(1).split("|")]
|
|
overrides = {}
|
|
for i, e in enumerate(exprs[:n_answers]):
|
|
val = safe_arithmetic(e)
|
|
if val is not None:
|
|
overrides[i] = str(int(val)) if float(val).is_integer() else str(val)
|
|
return overrides
|
|
# ---- Generation. enable_thinking=False keeps Qwen3 in its fast, non-<think>
|
|
# mode; our decomposition prompt supplies the reasoning instead. The kwarg is
|
|
# harmless on templates that ignore it. Both API-shape branches pass it. ----
|
|
def generate(messages, max_new_tokens, constraint_fn=None):
|
|
def _try_generate(gen_kwargs):
|
|
try:
|
|
enc = tok.apply_chat_template(
|
|
messages, add_generation_prompt=True, enable_thinking=False,
|
|
return_tensors="pt", return_dict=True,
|
|
).to(model.device)
|
|
input_len = enc["input_ids"].shape[-1]
|
|
with torch.no_grad():
|
|
out = model.generate(**enc, **gen_kwargs)
|
|
except Exception:
|
|
ids = tok.apply_chat_template(
|
|
messages, add_generation_prompt=True, enable_thinking=False,
|
|
return_tensors="pt",
|
|
).to(model.device)
|
|
input_len = ids.shape[-1]
|
|
with torch.no_grad():
|
|
out = model.generate(ids, **gen_kwargs)
|
|
return out, input_len
|
|
base_kwargs = {"max_new_tokens": max_new_tokens, "do_sample": False}
|
|
if constraint_fn is not None:
|
|
try:
|
|
out, input_len = _try_generate({**base_kwargs, "prefix_allowed_tokens_fn": constraint_fn})
|
|
except Exception:
|
|
out, input_len = _try_generate(base_kwargs)
|
|
else:
|
|
out, input_len = _try_generate(base_kwargs)
|
|
return tok.decode(out[0][input_len:], skip_special_tokens=True).strip()
|
|
EXPLANATION_SYSTEM = (
|
|
"Summarize the following reasoning into a few short bullet points: the "
|
|
"rule or pattern found in the data and the key evidence for the answer. "
|
|
"Be concise and structured -- do not repeat the full reasoning."
|
|
)
|
|
EXPLANATION_FALLBACK = "Answer derived from patterns found in the examples above."
|
|
_LETTER_CONSTRAINT_CACHE = {}
|
|
def build_letter_constraint_fn(tok, valid_letters):
|
|
cache_key = (id(tok), tuple(sorted(valid_letters)))
|
|
if cache_key in _LETTER_CONSTRAINT_CACHE:
|
|
return _LETTER_CONSTRAINT_CACHE[cache_key]
|
|
try:
|
|
allowed_chars = set(valid_letters) | set(" |\n\t\r")
|
|
eos = tok.eos_token_id
|
|
pieces = []
|
|
for token_id in range(len(tok)):
|
|
if token_id == eos:
|
|
continue
|
|
piece = tok.decode([token_id], skip_special_tokens=False)
|
|
if piece and all(c in allowed_chars for c in piece):
|
|
pieces.append(token_id)
|
|
allowed_ids = ([eos] if eos is not None else []) + pieces
|
|
def allowed(_batch_id, _input_ids):
|
|
return allowed_ids if allowed_ids else list(range(len(tok)))
|
|
_LETTER_CONSTRAINT_CACHE[cache_key] = allowed
|
|
return allowed
|
|
except Exception:
|
|
return None
|
|
def extract_match_letter_options(context: str):
|
|
found = set()
|
|
for line in context.splitlines():
|
|
for m in re.finditer(r"(?:^|\s)([A-Z])[.\)]\s+\S", line):
|
|
found.add(m.group(1))
|
|
if not found:
|
|
return None
|
|
letters = sorted(found)
|
|
expected = [chr(ord("A") + i) for i in range(len(letters))]
|
|
if letters != expected:
|
|
return None
|
|
if not (2 <= len(letters) <= 26):
|
|
return None
|
|
return letters
|
|
rows = []
|
|
processed_ids = set()
|
|
try:
|
|
for _, r in df.iterrows():
|
|
try:
|
|
elapsed = time.time() - start_time
|
|
remaining = TIME_LIMIT_S - elapsed
|
|
budget_left_rows = max(n_rows - len(rows), 1)
|
|
row_budget = remaining / budget_left_rows
|
|
time_based_cap = 1280 if row_budget > per_row_budget else 640
|
|
task_type = r.get("task_type", "")
|
|
messages, n_items = build_messages(r["context"], r["query"], task_type)
|
|
if n_items:
|
|
item_based_cap = max(640, min(1536, n_items * 48 + 256))
|
|
tokens_cap = min(time_based_cap, item_based_cap)
|
|
else:
|
|
tokens_cap = time_based_cap
|
|
text = generate(messages, tokens_cap)
|
|
answers, marker_pos = extract(text)
|
|
if task_type == "text_to_num":
|
|
overrides = extract_compute_overrides(text, len(answers))
|
|
for idx, val in overrides.items():
|
|
if idx < len(answers):
|
|
answers[idx] = val
|
|
if (marker_pos is None or not answers) and remaining > SETUP_BUFFER_S:
|
|
repair_constraint = None
|
|
if task_type == "match_letters":
|
|
repair_options = extract_match_letter_options(r["context"])
|
|
if repair_options:
|
|
repair_constraint = build_letter_constraint_fn(tok, repair_options)
|
|
repair_text = generate(build_repair_messages(r["query"], n_items, text), 128,
|
|
constraint_fn=repair_constraint)
|
|
rep, rep_pos = extract(repair_text)
|
|
if rep:
|
|
answers, marker_pos = rep, rep_pos
|
|
if n_items is not None:
|
|
if len(answers) < n_items:
|
|
answers = answers + [answers[-1] if answers else ""] * (n_items - len(answers))
|
|
elif len(answers) > n_items and marker_pos is None:
|
|
answers = answers[:n_items]
|
|
if not answers:
|
|
answers = [""]
|
|
remaining_after = TIME_LIMIT_S - (time.time() - start_time)
|
|
budget_left_after = max(n_rows - len(rows) - 1, 0)
|
|
comfortable = remaining_after > (budget_left_after + 1) * per_row_budget * 1.3
|
|
if comfortable:
|
|
try:
|
|
explanation = generate(
|
|
[{"role": "system", "content": EXPLANATION_SYSTEM},
|
|
{"role": "user", "content": text}], 300,
|
|
) or EXPLANATION_FALLBACK
|
|
except Exception:
|
|
explanation = EXPLANATION_FALLBACK
|
|
else:
|
|
snippet = re.sub(r"\s{2,}", " ", text[:300]).strip()
|
|
explanation = snippet if snippet else EXPLANATION_FALLBACK
|
|
rows.append({"id": r["id"], "pred": json.dumps(answers, ensure_ascii=False),
|
|
"explanation": explanation})
|
|
processed_ids.add(r["id"])
|
|
write_submission_csv(rows)
|
|
print(f"{len(rows)}/{n_rows} answers={len(answers)} elapsed={time.time()-start_time:.0f}s", flush=True)
|
|
except Exception as e:
|
|
try:
|
|
_, fallback_items, fk = parse_items(r["query"])
|
|
n_fallback = len(fallback_items) if fk else 1
|
|
except Exception:
|
|
n_fallback = 1
|
|
rows.append({"id": r["id"], "pred": json.dumps([""] * n_fallback, ensure_ascii=False),
|
|
"explanation": EXPLANATION_FALLBACK})
|
|
processed_ids.add(r["id"])
|
|
write_submission_csv(rows)
|
|
print(f"ROW ERROR on {r['id']}: {e}", flush=True)
|
|
if time.time() - start_time > TIME_LIMIT_S - 60:
|
|
print("Time budget nearly exhausted, stopping early.", flush=True)
|
|
break
|
|
for _, r in df.iterrows():
|
|
if r["id"] in processed_ids:
|
|
continue
|
|
try:
|
|
_, fallback_items, fk = parse_items(r["query"])
|
|
n_fallback = len(fallback_items) if fk else 1
|
|
except Exception:
|
|
n_fallback = 1
|
|
rows.append({"id": r["id"], "pred": json.dumps([""] * n_fallback, ensure_ascii=False),
|
|
"explanation": EXPLANATION_FALLBACK})
|
|
write_submission_csv(rows)
|
|
print("DONE.", flush=True)
|
|
except Exception as e:
|
|
emergency_submission_csv(f"main loop failed: {e}", rows_so_far=rows if rows else None)
|
|
print(f"FATAL, but submission.csv was written with {len(rows)} rows. Error: {e}", flush=True) |