rita-cohere/tya-m1-multilingual

1"""IOL-AI 2026 — M1 iterate on best private~0.083 (temp0.6 + user-instr).2 3Keep: user-prompt instructions, /think, sample think @ TEMPERATURE.4CSV fixes from that run: strip _GCY/gloss; reject essay+alphabet+resample;5stronger user prompt; think 2048; greedy answer continuation.6"""7 8import os9import subprocess10import sys11 12 13def _install_bundled_deps() -> None:14    wheels_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "wheels")15    if not os.path.isdir(wheels_dir):16        return17    subprocess.run(18        [19            sys.executable,20            "-m",21            "pip",22            "install",23            "-q",24            "--no-index",25            f"--find-links={wheels_dir}",26            "transformers==4.56.2",27        ],28        check=True,29    )30 31 32_install_bundled_deps()33 34os.environ["HF_HUB_OFFLINE"] = "1"35os.environ["TRANSFORMERS_OFFLINE"] = "1"36MODEL_ID = "."37 38# "" for A1 (reasoning_options only); "/think" for M1 multilingual39USER_THINK_TOKEN = "/think"40 41import json42import re43 44import pandas as pd45import torch46from transformers import AutoModelForCausalLM, AutoTokenizer47 48END_THINKING = "<|END_THINKING|>"49START_THINKING = "<|START_THINKING|>"50 51THINKING_BUDGET = 204852ANSWER_CONTINUATION_TOKENS = 51253COT_MAX_NEW_TOKENS = 102454TEMPERATURE = 0.655TOP_P = 0.9556QUALITY_RESAMPLES = 457ANSWER_GREEDY = True  # format-sensitive answer phase58 59# Empty system — instructions go on the user turn (decode ablation vs v4)60SYSTEM = ""61 62USER_INSTRUCTIONS = """You solve International Linguistics Olympiad (IOL) problems from the data you are given.63You may see a task type you have never seen: follow the instruction and examples, and answer in the same form they use.64 65What to return by task type:66- translation: only the required form in the language the query asks for — do not add extra glosses or "form | meaning" unless asked67- fill_blanks: only the missing form for each blank — no extra glosses68- match_letters: ONLY the option letter (A, B, C, …), one letter per line — never copy option text, never arrows, never "A. word"69- text_to_num: the number in digits only70- num_to_text: the number written out in words, in the language asked71- kinship / sentence matching: the full required sentence or form — NOT roman numerals (i, ii, iii) and NOT an alphabet dump72- any other type: exactly what the instruction asks for, nothing else73 74Answer in the language and form the query asks for. Do not add glosses, translations, or explanations unless the instruction requires them.75 76Output rules:77- Put answers ONLY after a line that says exactly: FINAL ANSWERS:78- Never put answers before that marker.79- One answer per line; exactly as many lines as items asked in the query.80- Bare answers only: no numbering, no quotes, no commentary, no repeating the question.81 82Extra hard rules:83- Never refuse or apologize; always output FINAL ANSWERS: with your best guess.84- For match_letters: only bare letters (A, B, C, …) — never dump the alphabet (A B C D E F…), never option text.85- Never append tags or glosses: no "_GCY", "_NS", "form – meaning", "word - gloss", or markdown bold.86- Never write an essay or explanation of how the language works under FINAL ANSWERS: — only the answer strings.87- If the query asks for colour/color forms, output those forms only (one per line), not a linguistics write-up.88- Emit exactly as many answer lines as items asked — no more, no fewer."""89 90USER_INSTRUCTIONS_COT = (91    USER_INSTRUCTIONS92    + "\n\nThink step by step about the rules in the examples and how they apply to the query, "93    "then write FINAL ANSWERS: and the answer lines."94)95 96# --- parser (inlined from parse_iol.py) ---97_MD_PREFIX = r"(?:[#*_=\-\s`>]*)"98_MARKER = re.compile(99    rf"(?im)^{_MD_PREFIX}final\s+answers?{_MD_PREFIX}:?{_MD_PREFIX}\s*(.*)$"100)101_NUMBERING = re.compile(r"^\s*(?:\d+[.)]|[-*•])\s*")102_TURN_NOISE = re.compile(103    r"<\|/?END_OF_TURN_TOKEN\|>|<\|/?START_OF_TURN_TOKEN\|>|"104    r"<\|CHATBOT_TOKEN\|>|<EOS_TOKEN>|<BOS_TOKEN>"105)106_RESPONSE_BLOCK = re.compile(107    r"<\|START_RESPONSE\|>(.*?)<\|END_RESPONSE\|>",108    flags=re.S,109)110_MD_WRAP = re.compile(r"^[*_`#\s]+|[*_`#\s]+$")111_TRAILING_LETTER = re.compile(112    r"(?:[–—\-]|→|->)\s*([A-Za-z])(?:\s*[.)]|)\s*$"113)114_LEADING_LETTER_OPT = re.compile(r"^([A-Za-z])\s*[.):\-–—]\s+\S")115_WORD_THEN_LETTER = re.compile(r"^.+\s([A-Za-z])\s*$")116_REFUSAL = re.compile(117    r"(?i)\b("118    r"i'?m sorry|i am sorry|i don'?t have|i cannot|i can'?t|"119    r"unable to|not able to|no reliable|cannot supply|can'?t supply|"120    r"as an ai|i apologize"121    r")\b"122)123 124 125def _strip_gloss_keep_form(line: str) -> str:126    s = (line or "").strip()127    s = re.sub(r"\*\*", "", s)128    s = re.split(r"\s+_?(?:GCY|NS|N/A)_?\b", s, maxsplit=1, flags=re.I)[0].strip()129    s = re.sub(r"\s+_?(?:GCY|NS|N/A)_?\s*$", "", s, flags=re.I).strip()130    m = re.match(131        r"^(.+?)\s+[-–—]\s+((?:to|the|a|an|in|of|for|being|means?|black|white|red|green|yellow)\b.*)$",132        s,133        flags=re.I,134    )135    if m:136        s = m.group(1).strip()137    return s.strip()138 139 140def _clean_line(line: str) -> str:141    line = _NUMBERING.sub("", line).strip()142    line = _MD_WRAP.sub("", line).strip()143    line = line.replace("\u202f", " ").replace("\xa0", " ")144    return _strip_gloss_keep_form(line.strip())145 146 147def after_thinking(text: str) -> str:148    """Prefer content after the last <|END_THINKING|>; else drop an unclosed think block."""149    if END_THINKING in text:150        text = text.rsplit(END_THINKING, 1)[-1]151    elif START_THINKING in text:152        text = ""153    return _TURN_NOISE.sub("", text)154 155 156def _as_option_letter(line: str) -> str | None:157    line = _clean_line(line)158    if not line:159        return None160    if len(line) == 1 and line.isalpha():161        return line.upper()162    m = _LEADING_LETTER_OPT.match(line)163    if m:164        return m.group(1).upper()165    m = _TRAILING_LETTER.search(line)166    if m:167        return m.group(1).upper()168    if len(line) <= 40:169        m = _WORD_THEN_LETTER.match(line)170        if m:171            return m.group(1).upper()172    return None173 174 175def _expand_line(line: str) -> list[str]:176    line = _clean_line(line)177    if not line:178        return []179    if len(line) == 1 and line.isalpha():180        return [line]181    if _LEADING_LETTER_OPT.match(line) or _TRAILING_LETTER.search(line):182        letter = _as_option_letter(line)183        if letter:184            return [letter]185    if len(line) <= 40 and _WORD_THEN_LETTER.match(line):186        letter = _as_option_letter(line)187        if letter:188            return [letter]189    if "|" in line:190        parts = [p.strip() for p in line.split("|") if p.strip()]191        if len(parts) >= 2:192            if len(parts) >= 4 and len(parts) % 2 == 0:193                left, right = parts[0::2], parts[1::2]194                if sum(" " in r for r in right) >= max(1, len(right) // 2):195                    return [_clean_line(x) for x in left if _clean_line(x)]196            if len(parts) == 2:197                a, b = parts198                if (" " in b and " " not in a) or (199                    len(b) > 2 * max(len(a), 1) and " " in b200                ):201                    return [_clean_line(a)] if _clean_line(a) else []202            return [_clean_line(p) for p in parts if _clean_line(p)]203    return [line]204 205 206def _dedupe_runaway(parts: list[str]) -> list[str]:207    if len(parts) < 6:208        return parts209    out: list[str] = []210    run = 0211    prev = None212    for p in parts:213        if p == prev:214            run += 1215            if run >= 4:216                break217        else:218            run = 1219            prev = p220        out.append(p)221    return out222 223 224def _lines_from_region(region: str, *, allow_all_lines: bool) -> list[str]:225    markers = list(_MARKER.finditer(region))226    if markers:227        last = markers[-1]228        after_parts: list[str] = []229        same = _clean_line(last.group(1) or "")230        if same:231            after_parts.extend(_expand_line(same))232        for line in region[last.end() :].splitlines():233            after_parts.extend(_expand_line(line))234        if after_parts:235            return _dedupe_runaway(after_parts)236        before_parts: list[str] = []237        for line in region[: last.start()].splitlines():238            before_parts.extend(_expand_line(line))239        if before_parts:240            return _dedupe_runaway(before_parts)241 242    parts: list[str] = []243    for line in region.splitlines():244        parts.extend(_expand_line(line))245    if not parts:246        return []247    if allow_all_lines:248        return _dedupe_runaway(parts)249    return [parts[-1]]250 251 252def parse_answers(253    raw: str,254    *,255    n_expected: int | None = None,256    task_type: str = "",257) -> list[str]:258    text = after_thinking(raw)259    closed_blocks = _RESPONSE_BLOCK.findall(text)260    answers: list[str] = []261    if closed_blocks:262        for region in reversed(closed_blocks):263            answers = _lines_from_region(region.strip(), allow_all_lines=True)264            if answers:265                break266    if not answers:267        answers = _lines_from_region(text, allow_all_lines=False)268 269    if task_type == "match_letters":270        coerced: list[str] = []271        for a in answers:272            letter = _as_option_letter(a)273            coerced.append(letter if letter else a)274        answers = coerced275 276    if n_expected is not None and n_expected > 0 and len(answers) > n_expected:277        answers = answers[:n_expected]278    return answers279 280 281def _looks_like_alphabet_dump(answers: list[str]) -> bool:282    letters = [a.strip().upper() for a in answers if len(a.strip()) == 1 and a.strip().isalpha()]283    if len(letters) < 8:284        return False285    seq = 0286    for i, L in enumerate(letters):287        if ord(L) == ord("A") + i:288            seq += 1289        else:290            break291    return seq >= 8292 293 294def _looks_like_essay(answers: list[str]) -> bool:295    if any(len(str(a)) > 100 for a in answers):296        return True297    blob = " ".join(map(str, answers))298    return bool(299        re.search(300            r"(?i)\b(colors? are expressed|systematic set of lexical|"301            r"these stems are combined|step by step|as an ai|verification)\b",302            blob,303        )304    )305 306 307def _looks_like_refusal(answers: list[str]) -> bool:308    blob = " ".join(answers)309    return bool(_REFUSAL.search(blob)) or len(blob) > 400 and "dictionary" in blob.lower()310 311 312def has_usable_answer(313    answers: list[str],314    *,315    n_expected: int | None = None,316    task_type: str = "",317) -> bool:318    if not answers or not any(a.strip() for a in answers):319        return False320    if _looks_like_refusal(answers):321        return False322    if _looks_like_alphabet_dump(answers):323        return False324    if _looks_like_essay(answers):325        return False326    if any(re.search(r"(?i)_GCY\b|_NS\b", str(a)) for a in answers):327        return False328    if task_type != "match_letters" and len(answers) >= 4:329        if all(re.fullmatch(r"[ivxlcdm]+", str(a).strip(), flags=re.I) for a in answers):330            return False331    if n_expected is not None and n_expected > 0 and abs(len(answers) - n_expected) > max(332        2, n_expected // 2333    ):334        return False335    if task_type == "match_letters":336        letters = [a for a in answers if len(a) == 1 and a.isalpha()]337        if len(letters) < max(1, int(0.8 * len(answers))):338            return False339    return True340 341 342def _n_items_guess(query: str) -> int:343    nums = re.findall(r"(?m)^\s*(?:\(?\d+[.)]|\d+\))", query)344    return len(nums) if nums else 0345 346 347def _end_thinking_id(tok) -> int:348    end_id = tok.convert_tokens_to_ids(END_THINKING)349    if end_id is None or end_id == tok.unk_token_id:350        ids = tok.encode(END_THINKING, add_special_tokens=False)351        if len(ids) == 1:352            end_id = ids[0]353    if end_id is None or end_id == tok.unk_token_id:354        raise RuntimeError(f"Tokenizer missing end-think token {END_THINKING!r}")355    return int(end_id)356 357 358def _build_prompt_ids(tok, system: str, user: str, *, thinking: bool):359    messages = []360    if system.strip():361        messages.append({"role": "system", "content": system})362    messages.append({"role": "user", "content": user})363    try:364        return tok.apply_chat_template(365            messages,366            add_generation_prompt=True,367            return_tensors="pt",368            reasoning_options={"enabled": thinking},369        )370    except TypeError:371        return tok.apply_chat_template(372            messages, add_generation_prompt=True, return_tensors="pt"373        )374 375 376def _sample_kwargs():377    return dict(378        do_sample=True,379        temperature=TEMPERATURE,380        top_p=TOP_P,381    )382 383 384@torch.inference_mode()385def generate_with_think_budget(model, tok, prompt_ids, end_id: int):386    device = next(model.parameters()).device387    prompt_ids = prompt_ids.to(device)388    prompt_len = prompt_ids.shape[-1]389    gen_kw = _sample_kwargs()390 391    think_out = model.generate(392        prompt_ids,393        max_new_tokens=THINKING_BUDGET,394        pad_token_id=tok.pad_token_id or tok.eos_token_id,395        **gen_kw,396    )[0]397    gen_ids = think_out[prompt_len:].tolist()398    if end_id not in gen_ids:399        cont = torch.cat(400            [think_out, torch.tensor([end_id], device=device, dtype=think_out.dtype)]401        )402    else:403        cont = think_out404 405    ans_kw = dict(do_sample=False) if ANSWER_GREEDY else gen_kw406    full = model.generate(407        cont.unsqueeze(0),408        max_new_tokens=ANSWER_CONTINUATION_TOKENS,409        pad_token_id=tok.pad_token_id or tok.eos_token_id,410        **ans_kw,411    )[0]412    text = tok.decode(full[prompt_len:], skip_special_tokens=False)413    return _TURN_NOISE.sub("", text).strip()414 415 416@torch.inference_mode()417def generate_plain(model, tok, prompt_ids, max_new_tokens: int):418    device = next(model.parameters()).device419    prompt_ids = prompt_ids.to(device)420    prompt_len = prompt_ids.shape[-1]421    out = model.generate(422        prompt_ids,423        max_new_tokens=max_new_tokens,424        pad_token_id=tok.pad_token_id or tok.eos_token_id,425        **_sample_kwargs(),426    )[0]427    text = tok.decode(out[prompt_len:], skip_special_tokens=False)428    return _TURN_NOISE.sub("", text).strip()429 430 431def _build_user(432    instructions: str,433    context: str,434    query: str,435    *,436    n_guess: int,437    think_token: str = "",438) -> str:439    parts = [440        instructions.strip(),441        "",442        context.strip(),443        "",444        query.strip(),445    ]446    if n_guess:447        parts.append("")448        parts.append(f"(Emit exactly {n_guess} answer line(s) after FINAL ANSWERS:.)")449    if think_token:450        parts.append(think_token.strip())451    return "\n".join(parts)452 453 454tok = AutoTokenizer.from_pretrained(MODEL_ID)455end_id = _end_thinking_id(tok)456model = AutoModelForCausalLM.from_pretrained(457    MODEL_ID, torch_dtype=torch.float16, device_map="auto"458).eval()459 460df = pd.read_csv("/tmp/data/test.csv", dtype=str).fillna("")461 462rows = []463for i, r in df.iterrows():464    n_guess = _n_items_guess(r["query"])465    task = str(r.get("task_type", "") or "")466    n_exp = n_guess or None467 468    user_think = _build_user(469        USER_INSTRUCTIONS,470        r["context"],471        r["query"],472        n_guess=n_guess,473        think_token=USER_THINK_TOKEN,474    )475    user_plain = _build_user(476        USER_INSTRUCTIONS_COT,477        r["context"],478        r["query"],479        n_guess=n_guess,480        think_token="",481    )482 483    best_answers: list[str] = []484    used_cot = False485    for attempt in range(QUALITY_RESAMPLES):486        torch.manual_seed(1000 + int(i) * 97 + attempt * 31)487        if torch.cuda.is_available():488            torch.cuda.manual_seed_all(1000 + int(i) * 97 + attempt * 31)489        ids = _build_prompt_ids(tok, SYSTEM, user_think, thinking=True)490        text = generate_with_think_budget(model, tok, ids, end_id)491        answers = parse_answers(text, n_expected=n_exp, task_type=task)492        answers = [_strip_gloss_keep_form(a) for a in answers if _strip_gloss_keep_form(a)]493        if n_exp and n_exp > 0 and len(answers) > n_exp:494            answers = answers[:n_exp]495        if has_usable_answer(answers, n_expected=n_exp, task_type=task):496            best_answers = answers497            break498        if answers and not best_answers:499            best_answers = answers500        print(501            f"  quality resample {attempt + 1}/{QUALITY_RESAMPLES} id={r['id']} n={len(answers)}",502            flush=True,503        )504    answers = best_answers505 506    if not has_usable_answer(answers, n_expected=n_exp, task_type=task):507        used_cot = True508        # CoT: no /think, thinking channel off, capped budget509        cot_ids = _build_prompt_ids(tok, SYSTEM, user_plain, thinking=False)510        cot_text = generate_plain(model, tok, cot_ids, COT_MAX_NEW_TOKENS)511        cot_answers = parse_answers(cot_text, n_expected=n_exp, task_type=task)512        cot_answers = [513            _strip_gloss_keep_form(a) for a in cot_answers if _strip_gloss_keep_form(a)514        ]515        if has_usable_answer(cot_answers, n_expected=n_exp, task_type=task) or (516            cot_answers and not answers517        ):518            answers = cot_answers519 520    rows.append({"id": r["id"], "pred": json.dumps(answers, ensure_ascii=False)})521    print(522        f"[{i + 1}/{len(df)}] {len(answers)} answers cot={used_cot}",523        flush=True,524    )525 526pd.DataFrame(rows).to_csv("submission.csv", index=False)527print("wrote submission.csv", flush=True)528