403 lines
14 KiB
Python
403 lines
14 KiB
Python
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 = 1200 # was 2000; cut to fit T4 wall clock
|
||
TEMPERATURE = 0.8
|
||
TOP_P = 0.95
|
||
MAX_ATTEMPTS = 1 # no per-lang retry; majority vote absorbs misses
|
||
NUM_LANGS = 3 # English + 2 others (fixed; was random 3–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 NUM_LANGS-1 others."""
|
||
others = [lang for lang in REASONING_LANGUAGES if lang != "English"]
|
||
return ["English"] + random.sample(others, NUM_LANGS - 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))
|
||
|
||
# Write incrementally so a wall-clock kill still leaves a partial submission.csv.
|
||
with open("submission.csv", "w", encoding="utf-8", newline="") as f:
|
||
writer = csv.DictWriter(f, fieldnames=["id", "pred"])
|
||
writer.writeheader()
|
||
f.flush()
|
||
|
||
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)
|
||
|
||
writer.writerow({"id": r["id"], "pred": json.dumps(voted, ensure_ascii=False)})
|
||
f.flush()
|
||
|
||
print("wrote submission.csv", flush=True)
|