| """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 = "user-seq2" |
|
|
| |
| os.environ["HF_HUB_OFFLINE"] = "1" |
| os.environ["TRANSFORMERS_OFFLINE"] = "1" |
| MODEL_ID = "." |
|
|
| VARIANT_CFG = { |
| |
| "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), |
| |
| "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), |
| |
| |
| "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 |
| 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 |
| ): |
| |
| 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": |
| |
| 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()) |
| |
| 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()) |
| ) |
| |
| 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 [] |
| |
| 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)) |
| |
| 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) |
|
|