79 lines
3.3 KiB
Python
79 lines
3.3 KiB
Python
import os
|
|
# 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 = "."
|
|
|
|
import re
|
|
import json
|
|
import pandas as pd
|
|
import torch
|
|
from transformers import AutoTokenizer, AutoModelForCausalLM
|
|
|
|
tok = AutoTokenizer.from_pretrained(MODEL_ID)
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
MODEL_ID, torch_dtype=torch.float16, device_map="auto"
|
|
).eval()
|
|
|
|
MAX_NEW_TOKENS = 1536
|
|
|
|
SYSTEM = (
|
|
"You solve International Linguistics Olympiad problems by reasoning from the "
|
|
"data you are given. You may meet a task type you have never seen: read the "
|
|
"instruction and the examples, and answer in the same form they use. "
|
|
"Common task types and what to give -- "
|
|
"translation: the translated form only, in the language the task asks for. "
|
|
"Apply every suffix, prefix, or ending shown in the examples (plurals, cases, "
|
|
"tense, etc.) -- do not give the bare stem if the pattern requires an ending; "
|
|
"fill_blanks: only the missing form for each blank; "
|
|
"match_letters: the option letter ALONE -- for example 'C', never 'word: C' or "
|
|
"any other text around it; "
|
|
"text_to_num: the number in digits; "
|
|
"num_to_text: the number written out in words, in the language asked; "
|
|
"any other type: give exactly what the instruction asks, nothing else. "
|
|
"Reason step by step first. Before you finalize, re-check each answer is in the "
|
|
"minimal bare form the task requires, with no echoed input, no labels, no extra "
|
|
"words attached. Then write a line that says exactly FINAL ANSWERS: "
|
|
"and, below it, one answer per line in the order the items are asked -- the "
|
|
"bare answer only, no numbering, no quotes, no extra text."
|
|
)
|
|
|
|
|
|
def parse_answers(text, task_type=None):
|
|
"""Keep only the lines after the last 'FINAL ANSWERS:' marker, one answer per line."""
|
|
marker = list(re.finditer(r"(?im)^\s*final answers?\s*:?\s*$", text))
|
|
if marker:
|
|
text = text[marker[-1].end():]
|
|
answers = []
|
|
for line in text.splitlines():
|
|
line = re.sub(r"^\s*\d+[.)]\s*", "", line).strip()
|
|
if line:
|
|
if task_type == "match_letters":
|
|
m = re.search(r"([A-Za-z])\s*$", line)
|
|
if m and (":" in line or len(line) > 3):
|
|
line = m.group(1)
|
|
answers.append(line)
|
|
return answers
|
|
|
|
|
|
df = pd.read_csv("/tmp/data/test.csv", dtype=str).fillna("")
|
|
|
|
rows = []
|
|
for _, r in df.iterrows():
|
|
messages = [
|
|
{"role": "system", "content": SYSTEM},
|
|
{"role": "user", "content": f"{r['context'].strip()}\n\n{r['query'].strip()}"},
|
|
]
|
|
enc = tok.apply_chat_template(
|
|
messages, add_generation_prompt=True, return_tensors="pt", return_dict=True,
|
|
).to(model.device)
|
|
with torch.no_grad():
|
|
out = model.generate(**enc, max_new_tokens=MAX_NEW_TOKENS, do_sample=False)
|
|
text = tok.decode(out[0][enc["input_ids"].shape[-1]:], skip_special_tokens=True).strip()
|
|
answers = parse_answers(text, task_type=r["task_type"])
|
|
rows.append({"id": r["id"], "pred": json.dumps(answers, ensure_ascii=False)})
|
|
print(f"{len(rows)}/{len(df)} done", flush=True)
|
|
|
|
os.makedirs("/tmp/model", exist_ok=True)
|
|
pd.DataFrame(rows).to_csv("/tmp/model/submission.csv", index=False)
|
|
print("wrote submission.csv", flush=True) |