script.py

# ============================================================
# 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)