140 lines
6.3 KiB
Python
140 lines
6.3 KiB
Python
import os
|
|
# The evaluation sandbox has NO internet. The model's weights are shipped inside
|
|
# this repo (at the repo root) and loaded from ".", with offline mode forced.
|
|
# transformers, torch, pandas, bitsandbytes and autoawq are already in the
|
|
# sandbox -- nothing can be pip installed at run time.
|
|
os.environ["HF_HUB_OFFLINE"] = "1"
|
|
os.environ["TRANSFORMERS_OFFLINE"] = "1"
|
|
|
|
import json, re, time
|
|
import pandas as pd
|
|
import torch
|
|
from transformers import AutoTokenizer, AutoModelForCausalLM
|
|
|
|
# The repo IS the working directory at run time -- weights live at the root,
|
|
# not in a subfolder. Do not nest the model in an inner folder.
|
|
MODEL_ID = "."
|
|
|
|
MAX_NEW_TOKENS = 1024 # 2560 risks timing out across a full test set in 30 min
|
|
TIME_BUDGET_SEC = 27 * 60 # safety margin under the 30-minute hard limit
|
|
|
|
print(f"Loading tokenizer from {MODEL_ID}...")
|
|
tok = AutoTokenizer.from_pretrained(MODEL_ID)
|
|
print("Tokenizer loaded.")
|
|
|
|
print(f"Loading model from {MODEL_ID}...")
|
|
# AWQ weights load automatically via autoawq (pre-installed) -- no special
|
|
# loader or install needed, from_pretrained detects quant_method from config.json.
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
MODEL_ID, torch_dtype=torch.float16, device_map="auto"
|
|
).eval()
|
|
print("Model loaded successfully.")
|
|
|
|
# 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("")
|
|
|
|
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."
|
|
)
|
|
|
|
_MARKER_RE = re.compile(r"(?im)^\s*[*_#\s]*final answers?[*_#\s]*:?\s*[*_#\s]*$")
|
|
_QUOTE_CHARS = "\"'`\u201c\u201d\u2018\u2019"
|
|
|
|
def _clean_line(line: str) -> str:
|
|
stripped = line.strip()
|
|
if re.match(r"^[\-\*\u2022]\s+", stripped):
|
|
stripped = re.sub(r"^[\-\*\u2022]\s+", "", stripped)
|
|
m = re.match(r"^\d+[.)]\s+", stripped)
|
|
if m:
|
|
stripped = stripped[m.end():]
|
|
return stripped.strip(_QUOTE_CHARS + " ")
|
|
|
|
def parse_answers(text):
|
|
marker = list(_MARKER_RE.finditer(text))
|
|
if marker:
|
|
text = text[marker[-1].end():]
|
|
|
|
lines = [ln for ln in text.splitlines() if ln.strip()]
|
|
lines = [ln for ln in lines if not re.match(r"^\s*```", ln)]
|
|
|
|
if len(lines) == 1 and "," in lines[0]:
|
|
parts = [p.strip().strip(_QUOTE_CHARS + " ") for p in lines[0].split(",")]
|
|
return [p for p in parts if p]
|
|
|
|
answers = [_clean_line(ln) for ln in lines]
|
|
return [a for a in answers if a]
|
|
|
|
def n_expected(query):
|
|
"""Count numbered items in the query so we can pad/trim a bad generation
|
|
instead of silently returning the wrong-length list (which zeroes that row)."""
|
|
nums = re.findall(r"(?m)^\s*(\d+)[.)]\s", query)
|
|
return len(nums) if nums else 1
|
|
|
|
def write_submission(rows):
|
|
# Save after every row -- guarantees a valid submission.csv exists even if
|
|
# the run gets cut off by the 30-minute limit.
|
|
pd.DataFrame(rows).to_csv("submission.csv", index=False)
|
|
|
|
rows = []
|
|
start = time.time()
|
|
for i, r in df.iterrows():
|
|
expected_n = n_expected(r["query"])
|
|
|
|
if time.time() - start > TIME_BUDGET_SEC:
|
|
print(f"Time budget hit at row {i+1}/{len(df)}, stopping early and saving progress.", flush=True)
|
|
for j in range(i, len(df)):
|
|
r2 = df.iloc[j]
|
|
rows.append({"id": r2["id"], "pred": json.dumps([""] * n_expected(r2["query"]), ensure_ascii=False)})
|
|
break
|
|
|
|
try:
|
|
messages = [
|
|
{"role": "system", "content": SYSTEM},
|
|
{"role": "user", "content": f"{r['context'].strip()}\n\n{r['query'].strip()}"},
|
|
]
|
|
ids = tok.apply_chat_template(
|
|
messages, add_generation_prompt=True, return_tensors="pt", return_dict=False
|
|
).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)
|
|
|
|
if len(answers) < expected_n:
|
|
answers = answers + [""] * (expected_n - len(answers))
|
|
elif len(answers) > expected_n:
|
|
answers = answers[:expected_n]
|
|
|
|
except Exception as e:
|
|
print(f"[{i+1}/{len(df)}] ERROR: {e}", flush=True)
|
|
answers = [""] * expected_n
|
|
|
|
rows.append({"id": r["id"], "pred": json.dumps(answers, ensure_ascii=False)})
|
|
print(f"[{i + 1}/{len(df)}] {len(answers)} answers", flush=True)
|
|
write_submission(rows)
|
|
|
|
write_submission(rows)
|
|
print("wrote submission.csv", flush=True) |