import os import subprocess import sys def _install_bundled_deps() -> None: """Install transformers from bundled wheels (eval sandbox has no PyPI access).""" 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 re import csv import json import random import shutil import tempfile import unicodedata from collections import Counter import torch from transformers import AutoTokenizer, AutoModelForCausalLM # 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 = "." MAX_NEW_TOKENS = 2000 TEMPERATURE = 0.8 TOP_P = 0.95 MAX_ATTEMPTS = 2 # per language; ensemble provides redundancy MIN_LANGS = 3 MAX_LANGS = 5 # Languages tiny-aya-global handles well for chain-of-thought. REASONING_LANGUAGES = [ "English", "Spanish", "French", "German", "Portuguese", "Italian", "Dutch", "Russian", "Arabic", "Simplified Chinese", "Japanese", "Korean", "Turkish", "Hindi", "Indonesian", "Vietnamese", "Polish", "Swedish", "Greek", "Hebrew", # "Swahili", "Ukrainian", "Romanian", "Czech", "Hungarian", ] def load_tokenizer(model_id: str = "."): """Load tokenizer, converting tokenizer.json for older tokenizers if needed.""" 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) # Older tokenizers expect merge pairs as "a b" strings, not ["a", "b"] lists. 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) SYSTEM_TEMPLATE = ( "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" "IMPORTANT: Write ALL of your step-by-step reasoning in {language}. " "Do not mix languages in the reasoning. " "The FINAL ANSWERS section must still use the English marker `FINAL ANSWERS:` " "and the answer values themselves must follow the TASK TYPE / QUERY requirements " "(do not translate those answers into {language} unless the query asks for that).\n\n" "As the first part of your answer, reason step by step in {language} 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." ) # Prefer a dedicated header line; also allow same-line answers after the colon. FINAL_ANSWERS_LINE_RE = re.compile( r"(?im)^[^\w\n]*final answers?[^\w\n]*:?[ \t]*(?=\n|$)|" r"(?im)^[^\w\n]*final answers?\s*:\s*" ) FINAL_ANSWERS_INLINE_RE = re.compile( r"(?is)\bfinal answers?\s*:\s*" ) def extract_raw_final(text: str) -> str: """Return text after the last final-answers marker, or '' if none found.""" line_matches = list(FINAL_ANSWERS_LINE_RE.finditer(text)) if line_matches: return text[line_matches[-1].end() :] inline_matches = list(FINAL_ANSWERS_INLINE_RE.finditer(text)) if inline_matches: return text[inline_matches[-1].end() :] return "" 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) -> list[str] | None: 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 parse_answer_lines(text_after_marker: str, query: str, task_type: str) -> list[str]: """Parse cleaned answer lines from the raw final-answers section.""" 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) if match_numbered_prefix: cleaned_line = match_numbered_prefix.group(1).strip() else: cleaned_line = 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 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 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) return answers def postprocess_answer(text, query, task_type): """Keep only the content after the last 'FINAL ANSWERS' marker.""" text_after_marker = extract_raw_final(text) if not text_after_marker.strip(): return [] return parse_answer_lines(text_after_marker, query, task_type) def normalize_for_vote(text: str, task_type: str) -> str: text = unicodedata.normalize("NFC", text.strip()) if task_type == "match_letters": return text.upper() return " ".join(text.split()) def majority_vote( lang_rollouts: list[tuple[str, list[str]]], expected: int, task_type: str, ) -> list[str]: """Per-item majority vote; ties break toward the English rollout.""" if expected <= 0: return [] # Prefer rollouts whose length matches the expected answer count. eligible = [ (lang, rollout) for lang, rollout in lang_rollouts if len(rollout) == expected and any(a.strip() for a in rollout) ] if not eligible: eligible = [ (lang, rollout) for lang, rollout in lang_rollouts if any(a.strip() for a in rollout) ] if not eligible: return [] final: list[str] = [] for i in range(expected): tagged = [ (lang, rollout[i]) for lang, rollout in eligible if i < len(rollout) and rollout[i].strip() ] if not tagged: final.append("") continue pairs = [ (lang, normalize_for_vote(ans, task_type), ans) for lang, ans in tagged ] counter = Counter(norm for _, norm, _ in pairs) top_count = max(counter.values()) tied_norms = {norm for norm, count in counter.items() if count == top_count} english_pair = next( ((norm, ans) for lang, norm, ans in pairs if lang == "English"), None, ) if english_pair is not None and english_pair[0] in tied_norms: winner_norm = english_pair[0] # Prefer English's surface form when it matches the winning norm. final.append(english_pair[1]) continue winner_norm = sorted(tied_norms)[0] originals = [ans for _, norm, ans in pairs if norm == winner_norm] final.append(Counter(originals).most_common(1)[0][0]) return final def sample_reasoning_languages() -> list[str]: """Always include English; sample the rest from the non-English pool.""" k = random.randint(MIN_LANGS, MAX_LANGS) others = [lang for lang in REASONING_LANGUAGES if lang != "English"] return ["English"] + random.sample(others, k - 1) def generate_for_language(tok, model, language: str, context: str, task_type: str, query: str) -> str: system = SYSTEM_TEMPLATE.format(language=language) messages = [ {"role": "system", "content": system}, { "role": "user", "content": ( f"CONTEXT:{context.strip()}\n" f"TASK TYPE:`{task_type}`\n\n" f"QUERY:{query.strip()}\n\n" f"Remember: reason entirely in {language}." ), }, ] ids = tok.apply_chat_template( messages, add_generation_prompt=True, return_tensors="pt", ).to(model.device) text = "" for attempt in range(1, MAX_ATTEMPTS + 1): with torch.no_grad(): out = model.generate( ids, max_new_tokens=MAX_NEW_TOKENS, do_sample=True, temperature=TEMPERATURE, top_p=TOP_P, ) text = tok.decode(out[0][ids.shape[-1] :], skip_special_tokens=True).strip() if extract_raw_final(text).strip(): return text print( f" [{language}] retry {attempt}/{MAX_ATTEMPTS}: no FINAL ANSWERS", flush=True, ) return text 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 = [] for idx, r in enumerate(test_rows, start=1): languages = sample_reasoning_languages() print(f"{idx}/{len(test_rows)} langs={languages}", flush=True) lang_rollouts: list[tuple[str, list[str]]] = [] for language in languages: text = generate_for_language( tok, model, language, r["context"], r["task_type"], r["query"], ) answers = postprocess_answer(text, r["query"], r["task_type"]) lang_rollouts.append((language, answers)) print(f" [{language}] parsed={answers!r}", flush=True) expected = expected_answer_count(r["query"], r["task_type"]) voted = majority_vote(lang_rollouts, expected, r["task_type"]) print(f" vote -> {voted!r}", flush=True) rows.append({"id": r["id"], "pred": json.dumps(voted, ensure_ascii=False)}) with open("submission.csv", "w", encoding="utf-8", newline="") as f: writer = csv.DictWriter(f, fieldnames=["id", "pred"]) writer.writeheader() writer.writerows(rows) print("wrote submission.csv", flush=True)