script.py

"""IOL Space script for Tiny Aya Global TTS variants.

Set VARIANT at the top before uploading to each Hub repo.
Julia P0 SYSTEM + original-style parser; T4-safe budgets per variant.
"""
import os
import subprocess
import sys


def _install_bundled_deps() -> None:
    wheels_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "wheels")
    if not os.path.isdir(wheels_dir):
        return
    subprocess.run(
        [
            sys.executable,
            "-m",
            "pip",
            "install",
            "-q",
            "--no-index",
            f"--find-links={wheels_dir}",
            "transformers==4.56.2",
        ],
        check=True,
    )


_install_bundled_deps()

import csv
import json
import random
import re
import shutil
import tempfile
from collections import Counter

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# ---------------------------------------------------------------------------
# VARIANT — overwritten by packaging script per Hub repo
# ---------------------------------------------------------------------------
VARIANT = "user-seq2"

# The repo is the working directory at run time, and there is no network.
os.environ["HF_HUB_OFFLINE"] = "1"
os.environ["TRANSFORMERS_OFFLINE"] = "1"
MODEL_ID = "."

VARIANT_CFG = {
    # user-prompt instructions + majvote k=3 + mild temp (T4-safe)
    "user-mv3": dict(k=3, temp=0.4, top_p=0.95, max_new=1200, attempts=1, mode="plain", soft_fmt=True, user_instr=True),
    # sequential: draft → revise, instructions on user turn
    "user-seq2": dict(k=1, temp=0.4, top_p=0.95, max_new=900, attempts=1, mode="iter2", soft_fmt=True, user_instr=True),
    "p0-reroll": dict(k=1, temp=0.8, top_p=0.95, max_new=2000, attempts=5, mode="plain", soft_fmt=False),
    "mv3": dict(k=3, temp=0.8, top_p=0.95, max_new=1200, attempts=2, mode="plain", soft_fmt=False),
    "mv3-tok1536": dict(k=3, temp=0.8, top_p=0.95, max_new=1536, attempts=1, mode="plain", soft_fmt=False),
    "mv5": dict(k=5, temp=0.8, top_p=0.95, max_new=900, attempts=1, mode="plain", soft_fmt=False),
    "greedy": dict(k=1, temp=0.0, top_p=1.0, max_new=1500, attempts=3, mode="greedy", soft_fmt=False),
    "t03": dict(k=1, temp=0.3, top_p=0.95, max_new=1500, attempts=3, mode="plain", soft_fmt=False),
    "induct": dict(k=1, temp=0.8, top_p=0.95, max_new=1024, attempts=2, mode="induct", soft_fmt=False, induct_tokens=768),
    "iter2": dict(k=1, temp=0.8, top_p=0.95, max_new=800, attempts=1, mode="iter2", soft_fmt=False),
    "analog": dict(k=1, temp=0.8, top_p=0.95, max_new=1200, attempts=2, mode="analog", soft_fmt=False, analog_tokens=600),
    "fmt-soft": dict(k=1, temp=0.8, top_p=0.95, max_new=1500, attempts=3, mode="plain", soft_fmt=True),
    # Julia P0 decode × 3 samples; pick by format (cardinality), not majority vote.
    # attempts=1 keeps T4 under ~30 min with max_new=2000.
    "fmtpick-k3": dict(
        k=3,
        temp=0.8,
        top_p=0.95,
        max_new=2000,
        attempts=1,
        mode="plain",
        soft_fmt=False,
        select="format",
    ),
}

CFG = VARIANT_CFG.get(VARIANT, VARIANT_CFG["p0-reroll"])
print(f"VARIANT={VARIANT} CFG={CFG}", flush=True)

SYSTEM = (
    "You solve International Linguistics Olympiad problems by reasoning from the "
    "data in CONTEXT you are given to solve the problems in QUERY. \n"
    "There are common TASK TYPES that we specify below, but "
    "you may meet a TASK TYPE you have never seen: read the "
    "instruction and the examples, and answer the QUERY in the same form they use.\n\n"
    "Common TASK TYPES and what to return: \n"
    "`translation`: return the translated form only, in the language the task asks for; \n"
    "`fill_blanks`: return only the missing form for each indicated blank "
    "(beware: this could be many different things: a word, a part of a word or a phonetic transcription---pay close attention to what part of the CONTEXT is missing in QUERY); \n"
    "`match_letters`: return only the option letter (for example A, B, C); \n"
    "`text_to_num`: return the number in digits; \n"
    "`num_to_text`: return the number written out in words, in the language asked; \n"
    "any other type: return exactly what the instruction asks for, nothing else. \n\n"
    "As the first part of your answer, reason step by step about (1) the linguistic "
    "rules that can be deduced from the given examples in CONTEXT, and (2) "
    "how to apply them to the given problems in QUERY, and (3) in what format answers need to be returned (words, numbers, phonetic transcriptions, ...). \n"
    "Then write a draft of the final answer. "
    "Subsequently, compare it with the format requirements again, "
    "and verify it's compliant with the deduced rules, and it is complete, i.e. has an answer for each element in QUERY. "
    "If necessary, correct and refine."
    "Finally, write a line that says exactly `FINAL ANSWERS:` "
    "and, below it, write the answers to the items requested in QUERY (not those in CONTEXT),"
    "one answer per line (separated by \n) in the order the items are asked for in the QUERY -- the "
    "bare answer only, no numbering, no quotes, no extra text, according to the given TASK TYPE."
)


USER_INSTRUCTIONS = SYSTEM  # same text; optionally placed on user turn
USE_USER_INSTR = bool(CFG.get("user_instr"))
SYSTEM_RUNTIME = "" if USE_USER_INSTR else SYSTEM

SYSTEM_INDUCT = (
    "You analyze International Linguistics Olympiad CONTEXT examples only. "
    "Deduce linguistic rules. Do NOT answer QUERY. "
    "Write a bullet list under a line that says exactly: RULES:"
)

SYSTEM_APPLY = (
    "You solve IOL problems using CONTEXT, RULES, TASK TYPE, and QUERY. "
    "Write a line `FINAL ANSWERS:` then one bare answer per line in QUERY order."
)

SYSTEM_ANALOG = (
    SYSTEM
    + "\n\nFirst invent 2 short analogical toy examples that mirror the rule pattern "
    "in CONTEXT (label them ANALOG:), then solve QUERY."
)

SYSTEM_REVISE = (
    "You revise draft answers for an IOL problem. Given CONTEXT, TASK TYPE, QUERY, "
    "and DRAFT, correct format and linguistic errors. "
    "Output only `FINAL ANSWERS:` then one bare answer per line."
)


def load_tokenizer(model_id: str = "."):
    tokenizer_path = os.path.join(model_id, "tokenizer.json")
    with open(tokenizer_path, encoding="utf-8") as handle:
        data = json.load(handle)
    merges = data.get("model", {}).get("merges", [])
    if not merges or not isinstance(merges[0], list):
        return AutoTokenizer.from_pretrained(model_id)
    data["model"]["merges"] = [" ".join(piece) for piece in merges]
    tmpdir = tempfile.mkdtemp()
    for name in ("tokenizer_config.json", "special_tokens_map.json"):
        src = os.path.join(model_id, name)
        if os.path.isfile(src):
            shutil.copy(src, tmpdir)
    with open(os.path.join(tmpdir, "tokenizer.json"), "w", encoding="utf-8") as handle:
        json.dump(data, handle)
    return AutoTokenizer.from_pretrained(tmpdir)


def expected_answer_count(query: str, task_type: str) -> int:
    if task_type == "match_letters":
        numbered = re.findall(r"^\s*\d+\.", query, re.MULTILINE)
        return len(numbered) or 1
    if "blanks" in query.lower():
        range_match = re.search(r"\((\d+)-(\d+)\)", query)
        if range_match:
            return int(range_match.group(2)) - int(range_match.group(1)) + 1
        return len(re.findall(r"\(\d+\)", query)) or 1
    numbered = re.findall(r"^\s*\d+[.)]", query, re.MULTILINE)
    return len(numbered) or 1


def split_single_line_answer(text: str, expected: int, task_type: str) -> list[str]:
    text = text.strip()
    if expected <= 1:
        return [text]

    def try_split(pattern: str):
        parts = [part.strip() for part in re.split(pattern, text) if part.strip()]
        return parts if len(parts) == expected else None

    if task_type == "match_letters":
        for pattern in (r"\s+", r",\s*", r";\s*"):
            if result := try_split(pattern):
                return result
        letters = re.findall(r"[A-Za-z]", text)
        if len(letters) == expected:
            return [letter.upper() for letter in letters]
        return [text]
    if task_type in ("text_to_num", "num_to_text"):
        for pattern in (r",\s*", r";\s*", r"\s+"):
            if result := try_split(pattern):
                return result
        return [text]
    for pattern in (r";\s*", r",\s*"):
        if result := try_split(pattern):
            return result
    return [text]


def postprocess_answer(text, query, task_type, soft_fmt: bool = False):
    marker_match = list(
        re.finditer(r"(?im)^[^\w\n]*final answers?[^\w\n]*:?\s*$", text)
    )
    if not marker_match:
        return []
    text_after_marker = text[marker_match[-1].end() :]
    answers = []
    for line in text_after_marker.splitlines():
        stripped_line = line.strip("`").strip()
        if stripped_line == "":
            continue
        match_numbered_prefix = re.match(r"^\s*\d+[.)]\s+(.*)", stripped_line)
        cleaned_line = (
            match_numbered_prefix.group(1).strip()
            if match_numbered_prefix
            else stripped_line
        )
        cleaned_line = re.sub(r"\*\*", "", cleaned_line).strip()

        if task_type == "match_letters":
            parts = [
                part.strip("().[]")
                for part in re.split(r"[\s,;]+", cleaned_line)
                if part.strip()
            ]
            if soft_fmt and len(parts) > 1 and all(
                re.fullmatch(r"[A-Za-z]", p) for p in parts
            ):
                # Soft: split letter-clusters into one answer per letter
                answers.extend([p.upper() for p in parts])
                continue
            if not (
                len(parts) > 1 and all(re.fullmatch(r"[A-Za-z]", part) for part in parts)
            ):
                match_letter_word = re.match(
                    r"^\s*(?:\(([A-Za-z])\)|\[([A-Za-z])\]|([A-Za-z]))\.?:?\s*(.*)$",
                    cleaned_line,
                )
                if match_letter_word:
                    letter = (
                        match_letter_word.group(1)
                        or match_letter_word.group(2)
                        or match_letter_word.group(3)
                    )
                    cleaned_line = letter.upper()

        if soft_fmt and task_type != "match_letters":
            # Strip "A: rest" only when rest is substantial
            m = re.match(
                r"^\s*(?:\(([A-Za-z])\)|\[([A-Za-z])\]|([A-Za-z]))\s*[.):\-–—]\s+(.+)$",
                cleaned_line,
            )
            if m and len(m.group(4).strip()) > 2:
                cleaned_line = m.group(4).strip()
            gloss = re.match(
                r"^(.+?)\s+[-–—]\s+((?:to|the|a|an|in|of|for|being)\b.*)$",
                cleaned_line,
                flags=re.I,
            )
            if gloss and sum(1 for c in gloss.group(2) if ord(c) > 127) <= 1:
                cleaned_line = gloss.group(1).strip()

        if cleaned_line:
            answers.append(cleaned_line)

    expected = expected_answer_count(query, task_type)
    if len(answers) == 1 and expected > 1:
        answers = split_single_line_answer(answers[0], expected, task_type)
    if soft_fmt and expected > 1 and len(answers) > expected:
        answers = answers[:expected]
    return answers


def majority_vote(samples: list[list[str]], expected: int) -> list[str]:
    if not samples:
        return []
    padded = [(list(s) + [""] * expected)[:expected] for s in samples]
    tup_counts = Counter(tuple(s) for s in padded)
    best_tup, best_c = tup_counts.most_common(1)[0]
    if best_c >= 2:
        return list(best_tup)
    out = []
    for i in range(expected):
        out.append(Counter(s[i] for s in padded).most_common(1)[0][0])
    return out


def _looks_like_gloss(text: str) -> bool:
    t = (text or "").strip()
    if not t:
        return False
    if re.search(r"\s+[-–—]\s+((?:to|the|a|an|in|of|for|being)\b)", t, flags=re.I):
        return True
    if len(t) > 40 and re.search(r"\b(means|meaning|gloss|translation)\b", t, flags=re.I):
        return True
    return False


def format_rank_key(pred: list[str], expected: int, task_type: str) -> tuple:
    """Higher tuple = better format fit (no gold labels)."""
    n = len(pred)
    exact_len = 1 if expected > 0 and n == expected else 0
    nonempty = sum(1 for p in pred if str(p).strip())
    # Prefer filling expected slots without empties / extras
    slots = (list(pred) + [""] * max(expected, 0))[: max(expected, 0)] if expected else list(pred)
    empty_slots = sum(1 for p in slots if not str(p).strip())
    length_pen = abs(n - expected) if expected else 0
    gloss_pen = sum(1 for p in pred if _looks_like_gloss(str(p)))
    letter_ok = 0
    if task_type == "match_letters" and slots:
        letter_ok = sum(
            1 for p in slots if re.fullmatch(r"[A-Za-z]", str(p).strip())
        )
    # Prefer exact cardinality, then filled slots, then letter shape, then fewer glosses
    return (
        exact_len,
        nonempty - empty_slots,
        letter_ok,
        -gloss_pen,
        -length_pen,
        nonempty,
    )


def pick_by_format(
    samples: list[list[str]], expected: int, task_type: str
) -> list[str]:
    """Choose the sample whose answer list best matches expected format."""
    if not samples:
        return []
    # Keep first occurrence on ties (stable)
    ranked = sorted(
        enumerate(samples),
        key=lambda pair: format_rank_key(pair[1], expected, task_type),
        reverse=True,
    )
    best_i, best = ranked[0]
    key = format_rank_key(best, expected, task_type)
    print(
        f"  fmtpick chose sample {best_i+1}/{len(samples)} "
        f"n_pred={len(best)} expected={expected} key={key}",
        flush=True,
    )
    return best


def chat_ids(tok, system: str, user: str):
    messages = []
    if (system or "").strip():
        messages.append({"role": "system", "content": system})
    messages.append({"role": "user", "content": user})
    return tok.apply_chat_template(
        messages, add_generation_prompt=True, return_tensors="pt"
    ).to(model.device)


def _with_user_instr(user: str) -> str:
    if not USE_USER_INSTR:
        return user
    return USER_INSTRUCTIONS.strip() + "\n\n" + user


@torch.inference_mode()
def generate_text(
    ids,
    *,
    max_new: int,
    temp: float,
    top_p: float,
    greedy: bool,
    seed: int,
    attempts: int,
    require_final: bool = True,
) -> str:
    text = ""
    for attempt in range(max(attempts, 1)):
        torch.manual_seed(seed + attempt * 9973)
        random.seed(seed + attempt * 9973)
        if torch.cuda.is_available():
            torch.cuda.manual_seed_all(seed + attempt * 9973)
        kwargs = dict(
            max_new_tokens=max_new,
            pad_token_id=tok.pad_token_id or tok.eos_token_id,
        )
        if greedy or temp <= 0:
            kwargs["do_sample"] = False
        else:
            kwargs.update(do_sample=True, temperature=temp, top_p=top_p)
        out = model.generate(ids, **kwargs)
        text = tok.decode(out[0][ids.shape[-1] :], skip_special_tokens=True).strip()
        if not require_final:
            return text
        low = text.lower()
        if "final answer" in low and low.split("final answer", 1)[1].strip():
            return text
        print(f"  retry {attempt+2}/{attempts}", flush=True)
    return text


def user_plain(r) -> str:
    base = (
        f"CONTEXT:{r['context'].strip()}\nTASK TYPE:`{r['task_type']}`\n\n"
        f"QUERY:{r['query'].strip()}"
    )
    return _with_user_instr(base)


tok = load_tokenizer(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID, torch_dtype=torch.float16, device_map="auto"
).eval()

with open("/tmp/data/test.csv", encoding="utf-8", newline="") as f:
    test_rows = list(csv.DictReader(f))

rows_out = []
for i, r in enumerate(test_rows):
    task_type = r["task_type"]
    query = r["query"]
    expected = expected_answer_count(query, task_type)
    mode = CFG["mode"]
    soft = CFG.get("soft_fmt", False)
    k = int(CFG["k"])
    samples_pred = []

    if mode == "induct":
        ids_r = chat_ids(
            tok,
            SYSTEM_INDUCT,
            f"CONTEXT:{r['context'].strip()}\nTASK TYPE:`{task_type}`\nDeduce RULES only.",
        )
        rules = generate_text(
            ids_r,
            max_new=int(CFG.get("induct_tokens", 768)),
            temp=CFG["temp"],
            top_p=CFG["top_p"],
            greedy=False,
            seed=1000 + i,
            attempts=1,
            require_final=False,
        )
        m = re.search(r"(?im)^rules:\s*", rules)
        if m:
            rules = rules[m.end() :].strip()
        ids_a = chat_ids(
            tok,
            SYSTEM_APPLY,
            f"CONTEXT:{r['context'].strip()}\nRULES:\n{rules}\n"
            f"TASK TYPE:`{task_type}`\n\nQUERY:{query.strip()}",
        )
        text = generate_text(
            ids_a,
            max_new=CFG["max_new"],
            temp=CFG["temp"],
            top_p=CFG["top_p"],
            greedy=False,
            seed=2000 + i,
            attempts=CFG["attempts"],
        )
        samples_pred.append(postprocess_answer(text, query, task_type, soft))

    elif mode == "analog":
        ids_an = chat_ids(tok, SYSTEM_ANALOG, user_plain(r))
        # short analog pass then solve in same generation
        text = generate_text(
            ids_an,
            max_new=CFG["max_new"],
            temp=CFG["temp"],
            top_p=CFG["top_p"],
            greedy=False,
            seed=3000 + i,
            attempts=CFG["attempts"],
        )
        samples_pred.append(postprocess_answer(text, query, task_type, soft))

    elif mode == "iter2":
        ids1 = chat_ids(tok, SYSTEM_RUNTIME, user_plain(r))
        draft = generate_text(
            ids1,
            max_new=CFG["max_new"],
            temp=CFG["temp"],
            top_p=CFG["top_p"],
            greedy=False,
            seed=4000 + i,
            attempts=1,
        )
        draft_ans = postprocess_answer(draft, query, task_type, soft)
        ids2 = chat_ids(
            tok,
            SYSTEM_REVISE,
            f"{user_plain(r)}\n\nDRAFT:\n" + "\n".join(draft_ans),
        )
        text = generate_text(
            ids2,
            max_new=CFG["max_new"],
            temp=CFG["temp"],
            top_p=CFG["top_p"],
            greedy=False,
            seed=4001 + i,
            attempts=1,
        )
        revised = postprocess_answer(text, query, task_type, soft)
        samples_pred.append(revised if revised else draft_ans)

    else:
        greedy = mode == "greedy"
        for s in range(k):
            ids = chat_ids(tok, SYSTEM_RUNTIME, user_plain(r))
            text = generate_text(
                ids,
                max_new=CFG["max_new"],
                temp=CFG["temp"],
                top_p=CFG["top_p"],
                greedy=greedy,
                seed=5000 + i * 97 + s,
                attempts=CFG["attempts"],
            )
            samples_pred.append(postprocess_answer(text, query, task_type, soft))

    select = CFG.get("select")
    if select == "format":
        pred = pick_by_format(samples_pred, expected, task_type)
    elif k > 1:
        pred = majority_vote(samples_pred, expected)
    else:
        pred = samples_pred[0] if samples_pred else []

    rows_out.append({"id": r["id"], "pred": json.dumps(pred, ensure_ascii=False)})
    print(
        f"{i+1}/{len(test_rows)} done variant={VARIANT} "
        f"n_pred={len(pred)} expected={expected} select={select or ('majority' if k > 1 else 'single')}",
        flush=True,
    )

with open("submission.csv", "w", encoding="utf-8", newline="") as f:
    writer = csv.DictWriter(f, fieldnames=["id", "pred"])
    writer.writeheader()
    writer.writerows(rows_out)
print("wrote submission.csv", flush=True)