353 lines
17 KiB
Python
353 lines
17 KiB
Python
# ============================================================
|
||
# IOL-AI 2026: ALGEBRAIC INDUCTION SOLVER (SUBMISSION SCRIPT)
|
||
# ============================================================
|
||
|
||
import os
|
||
# The evaluation sandbox has no internet access.
|
||
# These environment variables force Transformers to use local files only.
|
||
os.environ["HF_HUB_OFFLINE"] = "1"
|
||
os.environ["TRANSFORMERS_OFFLINE"] = "1"
|
||
os.environ["TOKENIZERS_PARALLELISM"] = "false"
|
||
|
||
import gc, re, time, torch, json
|
||
import pandas as pd
|
||
from collections import Counter
|
||
from transformers import AutoTokenizer, AutoModelForCausalLM
|
||
|
||
# ────────────────────────────────────────────────────
|
||
# 1. CONFIGURATION AND MODEL LOAD
|
||
# ────────────────────────────────────────────────────
|
||
try:
|
||
del model, tok
|
||
except NameError:
|
||
pass
|
||
|
||
gc.collect()
|
||
torch.cuda.empty_cache()
|
||
|
||
MODEL_ID = "." # Load weights directly from the repository
|
||
MAX_TOKEN_BUDGET = 2048
|
||
INDUCT_MAX_TOKENS = 800
|
||
MAX_ATTEMPTS = 3
|
||
GLOBAL_TIME_LIMIT = 1700 # 28.3 minutes (safe margin under 30 min limit)
|
||
SC_TASKS = frozenset({"match_letters", "fill_blanks"})
|
||
SC_K = 3
|
||
|
||
print("Loading tokenizer and model...", flush=True)
|
||
tok = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
|
||
model = AutoModelForCausalLM.from_pretrained(
|
||
MODEL_ID, dtype=torch.float16, device_map="auto", trust_remote_code=True
|
||
).eval()
|
||
print("Model loaded successfully.", flush=True)
|
||
|
||
tok.padding_side = "left"
|
||
if tok.pad_token_id is None:
|
||
tok.pad_token = tok.eos_token or tok.unk_token
|
||
|
||
# ────────────────────────────────────────────────────
|
||
# 2. EOS DISCOVERY
|
||
# ────────────────────────────────────────────────────
|
||
def get_eos_ids(tokenizer, model):
|
||
eos_ids = set()
|
||
if tokenizer.eos_token_id is not None:
|
||
eos_ids.add(int(tokenizer.eos_token_id))
|
||
meos = getattr(model.generation_config, "eos_token_id", None)
|
||
if meos:
|
||
if isinstance(meos, (list, tuple, set)): eos_ids.update(int(x) for x in meos)
|
||
else: eos_ids.add(int(meos))
|
||
return sorted(list(eos_ids))
|
||
|
||
EOS_IDS = get_eos_ids(tok, model)
|
||
EOS_SET = set(EOS_IDS)
|
||
model.generation_config.eos_token_id = EOS_IDS
|
||
model.generation_config.pad_token_id = tok.pad_token_id
|
||
|
||
# ────────────────────────────────────────────────────
|
||
# 3. DYNAMIC CONTEXT READERS
|
||
# ────────────────────────────────────────────────────
|
||
_IPA_HINT = re.compile(r"[\u0250-\u02AF\u02B0-\u02FF\u0300-\u036F\u1D00-\u1D7Føœæðθŋɣʔ]")
|
||
_ASKS_NON_PHONETIC = re.compile(r"(?i)translate\s+into\s+english|write\s+(it\s+)?in\s+the\s+[\w'\u2019-]+\s+orthography|in\s+the\s+regular\s+orthography")
|
||
_ASKS_TRANSCRIPTION = re.compile(r"(?i)\b(transcribe|transcription|phonetic(ally)?)\b")
|
||
|
||
def _bracketed_forms(text: str) -> list[str]:
|
||
out = []
|
||
for m in re.finditer(r"\[([^\[\]\n]{1,40})\]", text):
|
||
inner = m.group(1).strip()
|
||
if not inner or re.fullmatch(r"[\d\s,.\-]+", inner): continue
|
||
out.append(inner)
|
||
return out
|
||
|
||
def is_phonetic_task(context: str, query: str, min_forms: int = 3) -> bool:
|
||
if _ASKS_NON_PHONETIC.search(query): return False
|
||
if _bracketed_forms(query) and not _ASKS_TRANSCRIPTION.search(query): return False
|
||
forms = _bracketed_forms(context) + _bracketed_forms(query)
|
||
if len(forms) < min_forms: return False
|
||
phonetic_looking = sum(1 for f in forms if _IPA_HINT.search(f) or ":" in f)
|
||
return phonetic_looking >= max(2, len(forms) // 4)
|
||
|
||
def count_items(query: str) -> int:
|
||
n = len(re.findall(r"(?m)^\s*\d+[.)]", query))
|
||
if n: return n
|
||
if "blanks" in query.lower():
|
||
m = re.search(r"\((\d+)-(\d+)\)", query)
|
||
if m: return int(m.group(2)) - int(m.group(1)) + 1
|
||
return len(re.findall(r"\(\d+\)", query)) or 0
|
||
return 0
|
||
|
||
# ────────────────────────────────────────────────────
|
||
# 4. DEFENSIVE PARSING
|
||
# ────────────────────────────────────────────────────
|
||
_TURN_NOISE = re.compile(r"<\|/?END_OF_TURN_TOKEN\|>|<\|/?START_OF_TURN_TOKEN\|>|<\|CHATBOT_TOKEN\|>|<EOS_TOKEN>|<BOS_TOKEN>|<\|im_end\|>|<\|im_start\|>")
|
||
_MARKER = re.compile(r"(?im)^\s*final answers?\s*:?\s*$")
|
||
|
||
def _looks_like_prose(line: str) -> bool:
|
||
if re.search(r"(?i)^(final answers?|answers?|note|reviewing|summary|explanation|verification)\b.*:$", line): return True
|
||
if re.search(r"(?i)^(here (are|is)|the (final )?answers? (are|is)|based on|therefore|thus|in summary)\b", line): return True
|
||
if line.rstrip().endswith(":") and len(line) > 3: return True
|
||
if len(line) > 120: return True
|
||
return False
|
||
|
||
def _strip_gloss_keep_form(line: str) -> str:
|
||
s = re.sub(r"\*\*", "", (line or "").strip())
|
||
s = re.split(r"\s+_?(?:GCY|NS|N/A)_?\b", s, maxsplit=1, flags=re.I)[0].strip()
|
||
m = re.match(r"^(.+?)\s+[-–—]\s+((?:to|the|a|an|in|of|for|being|means?)\b.*)$", s, flags=re.I)
|
||
if m: s = m.group(1).strip()
|
||
return s.strip()
|
||
|
||
def parse_answers(text: str, n_items: int = 0) -> list[str]:
|
||
text = _TURN_NOISE.sub("", text or "")
|
||
markers = list(_MARKER.finditer(text))
|
||
if markers:
|
||
text = text[markers[-1].end():]
|
||
else:
|
||
m = re.search(r'<final_answers>(.*?)</final_answers>', text, re.DOTALL | re.IGNORECASE)
|
||
if m: text = m.group(1)
|
||
|
||
answers = []
|
||
for line in text.splitlines():
|
||
line = re.sub(r"^\s*\d+[.)]\s*", "", line).strip().strip("`").strip("*")
|
||
if not line or _looks_like_prose(line): continue
|
||
|
||
line = _strip_gloss_keep_form(line)
|
||
if not line: continue
|
||
|
||
if re.fullmatch(r"(?:[A-Za-z]\s+)+[A-Za-z]", line):
|
||
answers.extend([p.upper() for p in line.split()])
|
||
continue
|
||
answers.append(line)
|
||
|
||
if n_items > 0:
|
||
answers = answers[:n_items]
|
||
if len(answers) < n_items:
|
||
answers += [""] * (n_items - len(answers))
|
||
return answers
|
||
|
||
def majority_vote(samples: list[list[str]], n_items: int) -> list[str]:
|
||
usable = [s for s in samples if any(x.strip() for x in s)]
|
||
if not usable: return [""] * max(n_items, 0)
|
||
n = n_items or max(len(s) for s in usable)
|
||
padded = [(list(s) + [""] * n)[:n] for s in usable]
|
||
counts = Counter(tuple(p) for p in padded)
|
||
best, c = counts.most_common(1)[0]
|
||
if c >= 2: return list(best)
|
||
return [Counter(p[i] for p in padded).most_common(1)[0][0] for i in range(n)]
|
||
|
||
# ────────────────────────────────────────────────────
|
||
# 5. PROMPT BUILDERS (ALGEBRAIC INDUCTION)
|
||
# ────────────────────────────────────────────────────
|
||
SYSTEM_BASE = (
|
||
"You are an elite computational linguist solving International Linguistics Olympiad problems. "
|
||
"This is a closed-world puzzle. DO NOT use your knowledge of real-world languages. "
|
||
"You may meet a task type you have never seen: read the instruction and the examples, and answer in the same form they use. "
|
||
"You MUST output your reasoning inside <reasoning> tags first. "
|
||
"After your reasoning is complete, you MUST write a line that says exactly FINAL ANSWERS: and, below it, "
|
||
"one answer per line in the order the items are asked -- the bare answer only, no numbering, "
|
||
"no quotes, no extra text. After FINAL ANSWERS:, output only the answers, exactly one line per "
|
||
"numbered item, then stop."
|
||
)
|
||
|
||
PHONETIC_INSTRUCTION = (
|
||
"IMPORTANT -- this problem uses PHONETIC TRANSCRIPTION. The examples write forms "
|
||
"inside square brackets, like [bø:va]. Your answers must be phonetic transcriptions "
|
||
"in exactly that same notation: enclosed in square brackets, using the same phonetic "
|
||
"symbols. Do NOT give an English meaning or gloss -- give the transcribed FORM."
|
||
)
|
||
|
||
def build_system(task_type: str, context: str, query: str) -> str:
|
||
parts = [SYSTEM_BASE]
|
||
if is_phonetic_task(context, query):
|
||
parts.append(PHONETIC_INSTRUCTION)
|
||
task_type = str(task_type).strip().lower()
|
||
if task_type == "match_letters":
|
||
parts.append("This is a MATCHING task. Answer with a SINGLE OPTION LETTER only (e.g., C).")
|
||
elif task_type == "text_to_num":
|
||
parts.append("This is a TEXT-TO-NUMBER task. Give the number in digits only (e.g., 111).")
|
||
return "\n\n".join(parts)
|
||
|
||
def build_user(row, n_items: int, rules: str = "", mode: str = "answer", error_feedback: str = None) -> str:
|
||
content = f"{str(row['context']).strip()}\n\n{str(row['query']).strip()}"
|
||
|
||
if mode == "induct":
|
||
content += (
|
||
"\n\nDeduce the linguistic system as a strict ALGEBRAIC EQUATION SHEET. DO NOT write prose. DO NOT answer the QUERY yet. "
|
||
"Inside <reasoning> tags, output ONLY the following mathematical notations based on the CONTEXT:\n\n"
|
||
"1. ALIGNMENT: Define the abstract structure using variables.\n"
|
||
" - If concatenative: `Word = A + B + C` (e.g., `anguls = angul + s`)\n"
|
||
" - If infixing: `Word = A + Infix + B` (e.g., `sumulat = s + um + ulat`)\n"
|
||
" - If templatic/ablaut: `Word = F(Root)` (e.g., `sang = Past(sing)`, `kataba = CaCaCa(k,t,b)`)\n"
|
||
" - If reduplication: `Word = A + A` (e.g., `bukubuku = buku + buku`)\n"
|
||
"2. MORPHOLOGY: Map variables to meanings. (e.g., `A = sing`, `Past = F()`, `s = Plural`)\n"
|
||
"3. PHONOLOGY: Write exact sound changes using rule notation: /input/ -> [output] / environment. (e.g., `/v/ -> [g] / ø:_a`)\n"
|
||
"4. MATCHING (if applicable): Map forms to options using matrices. (e.g., `u'u = breast = Option A`)\n"
|
||
"5. NUMBERS (if applicable): Map bases mathematically. (e.g., `123 = 6 * 20^1 + 3 * 20^0`)\n\n"
|
||
"Then write a line that says exactly: RULES:"
|
||
)
|
||
return content
|
||
|
||
if rules.strip():
|
||
content += (
|
||
f"\n\nINDUCED RULES:\n{rules.strip()}\n\n"
|
||
"CRITICAL: You must solve the algebraic equations from the RULES to construct the answers. "
|
||
"Do NOT guess. Do NOT blindly copy and paste full words from the context. "
|
||
"Apply the exact functions, morpheme slots, and sound changes to derive the final forms."
|
||
)
|
||
|
||
if n_items > 0:
|
||
content += f"\n\nThere are exactly {n_items} items to answer. Give exactly {n_items} answers after FINAL ANSWERS:, one per line, no more and no fewer."
|
||
|
||
if error_feedback:
|
||
content += f"\n\nPREVIOUS ATTEMPT FAILED:\n{error_feedback}\n\nFix your equation solving and output the corrected answers again."
|
||
|
||
return content
|
||
|
||
def extract_rules(text: str) -> str:
|
||
text = _TURN_NOISE.sub("", text or "")
|
||
m = list(re.finditer(r"(?im)^\s*rules?\s*:?\s*$", text))
|
||
if m: return text[m[-1].end():].strip()[:2000]
|
||
return text.strip()[:2000]
|
||
|
||
# ────────────────────────────────────────────────────
|
||
# 6. INFERENCE ENGINE
|
||
# ────────────────────────────────────────────────────
|
||
@torch.inference_mode()
|
||
def generate(prompt_text: str, max_new_tokens: int, sample: bool = False, seed: int = 0):
|
||
torch.manual_seed(seed)
|
||
if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)
|
||
|
||
enc = tok(prompt_text, return_tensors="pt", add_special_tokens=False).to(model.device)
|
||
plen = enc["input_ids"].shape[1]
|
||
|
||
kw = dict(do_sample=True, temperature=0.6, top_p=0.95) if sample else dict(do_sample=False)
|
||
|
||
out = model.generate(
|
||
**enc, max_new_tokens=max_new_tokens, use_cache=True,
|
||
eos_token_id=EOS_IDS, pad_token_id=tok.pad_token_id, **kw
|
||
)
|
||
gid = out[0, plen:]
|
||
|
||
for pos, tid in enumerate(gid.tolist()):
|
||
if tid in EOS_SET:
|
||
rt = tok.decode(gid[:pos+1], skip_special_tokens=False).strip()
|
||
return rt, pos+1
|
||
rt = tok.decode(gid, skip_special_tokens=False).strip()
|
||
return rt, gid.shape[0]
|
||
|
||
# ────────────────────────────────────────────────────
|
||
# 7. LOAD DATA & DYNAMIC INFERENCE LOOP
|
||
# ────────────────────────────────────────────────────
|
||
print("Loading test data...", flush=True)
|
||
df = pd.read_csv("/tmp/data/test.csv", dtype=str).fillna("")
|
||
|
||
if "id" not in df.columns:
|
||
df["id"] = df.index
|
||
|
||
results_by_id = {}
|
||
run_start = time.time()
|
||
order = df.index.tolist()
|
||
order.sort(key=lambda idx: len(str(df.loc[idx,"context"]))+len(str(df.loc[idx,"query"])))
|
||
|
||
current_token_budget = MAX_TOKEN_BUDGET
|
||
|
||
for i, row_id in enumerate(order):
|
||
row = df.loc[row_id]
|
||
task_type = str(row.get("task_type", "general"))
|
||
n_items = count_items(str(row["query"]))
|
||
|
||
t0 = time.time()
|
||
messages = []
|
||
final_answers = []
|
||
raw_output = ""
|
||
status = "FAIL"
|
||
gen_len = 0
|
||
|
||
elapsed = time.time() - run_start
|
||
remaining_time = GLOBAL_TIME_LIMIT - elapsed
|
||
problems_left = len(order) - i
|
||
max_allowed_tokens_for_time = max(256, int((remaining_time - (problems_left * 5)) * 10))
|
||
if max_allowed_tokens_for_time < current_token_budget:
|
||
current_token_budget = max_allowed_tokens_for_time
|
||
|
||
# PASS 1: Induction (Greedy, strict algebraic extraction)
|
||
system_prompt = build_system(task_type, str(row["context"]), str(row["query"]))
|
||
induct_user = build_user(row, n_items=0, mode="induct")
|
||
induct_prompt = tok.apply_chat_template(
|
||
[{"role":"system","content":system_prompt}, {"role":"user","content":induct_user}],
|
||
add_generation_prompt=True, tokenize=False
|
||
)
|
||
induct_raw, _ = generate(induct_prompt, INDUCT_MAX_TOKENS, sample=False, seed=1)
|
||
rules = extract_rules(induct_raw)
|
||
|
||
for attempt in range(MAX_ATTEMPTS):
|
||
error_feedback = messages[-1] if messages else None
|
||
user_prompt = build_user(row, n_items=n_items, rules=rules, mode="answer", error_feedback=error_feedback)
|
||
prompt = tok.apply_chat_template(
|
||
[{"role":"system","content":system_prompt}, {"role":"user","content":user_prompt}],
|
||
add_generation_prompt=True, tokenize=False
|
||
)
|
||
|
||
if task_type in SC_TASKS and attempt == 0:
|
||
samples = []
|
||
raw_samples = []
|
||
for k in range(SC_K):
|
||
raw_out, gen_len = generate(prompt, current_token_budget, sample=True, seed=1000+k*17)
|
||
raw_samples.append(raw_out)
|
||
samples.append(parse_answers(raw_out, n_items=n_items))
|
||
final_answers = majority_vote(samples, n_items)
|
||
raw_output = "\n---\n".join(raw_samples)
|
||
else:
|
||
raw_output, gen_len = generate(prompt, current_token_budget, sample=False, seed=42+attempt)
|
||
final_answers = parse_answers(raw_output, n_items=n_items)
|
||
|
||
if len(final_answers) < n_items or not all(final_answers):
|
||
msg = f"PARSE ERROR: Expected {n_items} answers, but extracted {len([a for a in final_answers if a])}. Ensure you output exactly {n_items} answers inside FINAL ANSWERS:."
|
||
messages.append(msg)
|
||
print(f"id={row_id} Attempt {attempt+1}: DENIED - {msg[:80]}", flush=True)
|
||
continue
|
||
|
||
status = "OK"
|
||
print(f"id={row_id} Attempt {attempt+1}: SUCCESS", flush=True)
|
||
break
|
||
|
||
if status != "OK":
|
||
final_answers = [""] * n_items if n_items > 0 else []
|
||
print(f"id={row_id} Failed after {MAX_ATTEMPTS} attempts.", flush=True)
|
||
|
||
wt = time.time() - t0
|
||
results_by_id[row_id] = {
|
||
"id": row["id"],
|
||
"pred": final_answers
|
||
}
|
||
print(f" tok={gen_len:>4}/{current_token_budget} time={wt:.1f}s total={int(elapsed)}s", flush=True)
|
||
|
||
# ────────────────────────────────────────────────────
|
||
# 8. WRITE SUBMISSION
|
||
# ────────────────────────────────────────────────────
|
||
out_rows = []
|
||
for res in results_by_id.values():
|
||
out_rows.append({
|
||
"id": res["id"],
|
||
"pred": json.dumps(res["pred"], ensure_ascii=False)
|
||
})
|
||
|
||
pd.DataFrame(out_rows, columns=["id", "pred"]).to_csv("submission.csv", index=False)
|
||
print("wrote submission.csv", flush=True) |