170 lines
7.6 KiB
Python
170 lines
7.6 KiB
Python
import os
|
|
# The evaluation sandbox has NO internet. The model's weights are shipped inside
|
|
# this repo and loaded from the local folder ".", with offline mode forced. We do
|
|
# NOT pip install anything: transformers, torch and pandas are already in the
|
|
# sandbox, and autoawq (needed for AWQ models) is preinstalled too.
|
|
os.environ["HF_HUB_OFFLINE"] = "1"
|
|
os.environ["TRANSFORMERS_OFFLINE"] = "1"
|
|
|
|
import json, re, shutil, tempfile
|
|
import pandas as pd
|
|
import torch
|
|
from transformers import AutoTokenizer, AutoModelForCausalLM
|
|
|
|
MODEL_ID = "." # the model's weights ship inside this repo
|
|
MAX_NEW_TOKENS = 2048 # this is a 1.7B model, cheap to run -- room for a longer answer block
|
|
|
|
|
|
def load_tokenizer(model_id: str = "."):
|
|
"""Load the tokenizer, patching tokenizer.json if it uses a merges format
|
|
the sandbox's (older) tokenizers build can't parse. Newer exports sometimes
|
|
store BPE merges as [["a","b"], ...] (list-of-lists) instead of the older
|
|
["a b", ...] (space-joined strings), which raises:
|
|
'data did not match any variant of untagged enum ModelWrapper'."""
|
|
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)
|
|
|
|
|
|
# 1) Load the model shipped in this repo (float16 = the T4's native precision).
|
|
tok = load_tokenizer(MODEL_ID)
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
MODEL_ID, torch_dtype=torch.float16, device_map="auto"
|
|
).eval()
|
|
|
|
# 2) Read the hidden test set the platform mounts for us (one row per problem).
|
|
df = pd.read_csv("/tmp/data/test.csv", dtype=str).fillna("")
|
|
|
|
# 3) How we ask: let the model reason, then write its answers after a marker.
|
|
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 "
|
|
"(this could be a word, 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"
|
|
"First, reason step by step about (1) the linguistic rules that can be deduced "
|
|
"from the examples in CONTEXT, and (2) how to apply them to the items in QUERY. "
|
|
"Then write a draft answer, check it against the format requirements and the "
|
|
"deduced rules, and make sure it has one answer for every item in QUERY. Correct "
|
|
"it if needed.\n"
|
|
"Finally, write a line that says exactly FINAL ANSWERS: and, below it, the "
|
|
"answers to the items in QUERY (not those already given in CONTEXT), one "
|
|
"answer per line, in the order the items are asked for -- the bare answer "
|
|
"only, no numbering, no quotes, no extra text."
|
|
)
|
|
|
|
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, expected, task_type):
|
|
text = text.strip()
|
|
if expected <= 1:
|
|
return [text]
|
|
|
|
def try_split(pattern):
|
|
parts = [p.strip() for p in re.split(pattern, text) if p.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]
|
|
|
|
for pattern in (r";\s*", r",\s*", r"\s+"):
|
|
if result := try_split(pattern):
|
|
return result
|
|
return [text]
|
|
|
|
def parse_answers(text, query, task_type):
|
|
"""Keep only the lines after the last 'FINAL ANSWERS:' marker, one per line.
|
|
We drop the reasoning above it and return the answers in order; the scorer
|
|
lines our list up against the reference by position."""
|
|
marker = list(re.finditer(r"(?im)^[^\w\n]*final answers?[^\w\n]*:?\s*$", text))
|
|
if not marker:
|
|
return []
|
|
text = text[marker[-1].end():]
|
|
|
|
answers = []
|
|
for line in text.splitlines():
|
|
line = line.strip("`").strip()
|
|
if not line:
|
|
continue
|
|
numbered = re.match(r"^\s*\d+[.)]\s+(.*)", line)
|
|
line = numbered.group(1).strip() if numbered else line
|
|
line = re.sub(r"\*\*", "", line).strip()
|
|
|
|
if task_type == "match_letters":
|
|
parts = [p.strip("().[]") for p in re.split(r"[\s,;]+", line) if p.strip()]
|
|
if not (len(parts) > 1 and all(re.fullmatch(r"[A-Za-z]", p) for p in parts)):
|
|
m = re.match(r"^\s*(?:\(([A-Za-z])\)|\[([A-Za-z])\]|([A-Za-z]))\.?:?\s*(.*)$", line)
|
|
if m:
|
|
line = (m.group(1) or m.group(2) or m.group(3)).upper()
|
|
|
|
if line:
|
|
answers.append(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
|
|
|
|
# 4) Answer every problem, in order, and write the submission file.
|
|
rows = []
|
|
for i, r in df.iterrows():
|
|
messages = [
|
|
{"role": "system", "content": SYSTEM},
|
|
{"role": "user", "content": (
|
|
f"CONTEXT:\n{r['context'].strip()}\n\n"
|
|
f"TASK TYPE: `{r['task_type']}`\n\n"
|
|
f"QUERY:\n{r['query'].strip()}"
|
|
)},
|
|
]
|
|
ids = tok.apply_chat_template(
|
|
messages, add_generation_prompt=True, return_tensors="pt",
|
|
enable_thinking=False, # Qwen3 supports a <think> mode; off keeps output short and predictable
|
|
).to(model.device)
|
|
with torch.no_grad():
|
|
out = model.generate(ids, max_new_tokens=MAX_NEW_TOKENS, do_sample=False)
|
|
text = tok.decode(out[0][ids.shape[-1]:], skip_special_tokens=True).strip()
|
|
answers = parse_answers(text, r["query"], r["task_type"])
|
|
rows.append({"id": r["id"], "pred": json.dumps(answers, ensure_ascii=False)})
|
|
print(f"[{i + 1}/{len(df)}] {len(answers)} answers", flush=True)
|
|
|
|
pd.DataFrame(rows).to_csv("submission.csv", index=False)
|