212 lines
8.6 KiB
Python
212 lines
8.6 KiB
Python
import os
|
|
os.environ["HF_HUB_OFFLINE"] = "1"
|
|
os.environ["TRANSFORMERS_OFFLINE"] = "1"
|
|
|
|
import csv, json, re, time
|
|
from collections import Counter
|
|
import torch
|
|
from transformers import AutoTokenizer, AutoModelForCausalLM
|
|
|
|
MODEL_ID = "." # weights ship inside this repo
|
|
MAX_NEW_TOKENS = 1536
|
|
K_SAMPLES = 3
|
|
TEMPERATURE = 0.7
|
|
WRITE_EXPLANATIONS = True
|
|
TIME_LIMIT_SECONDS = 1500
|
|
|
|
SYSTEM = (
|
|
"You solve International Linguistic Olympiad problems by reasoning from the "
|
|
"data given in CONTEXT. Everything you need is there; you don't need to know "
|
|
"the language in advance. Work step by step: find the pattern that links the "
|
|
"paired forms, verify it against EVERY example in CONTEXT (not just one), then "
|
|
"apply it to each item in QUERY. Before finalizing, sanity-check each answer "
|
|
"against the rule you found. You may meet a task type you've never seen -- read "
|
|
"the instruction and examples and answer in the same form they use. "
|
|
"Always give your best guess for every item, even if unsure -- partial credit "
|
|
"is awarded, a blank never is. "
|
|
"Reason step by step first. 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."
|
|
)
|
|
|
|
SUMMARIZE = (
|
|
"Summarize the following reasoning into a few short bullet points: the rule "
|
|
"or pattern found in the data and the key evidence for the answer. Be concise "
|
|
"and structured -- do not repeat the full reasoning."
|
|
)
|
|
|
|
FEWSHOT_EXAMPLES = {
|
|
"translation": {
|
|
"context": (
|
|
"Here are some sentences in Hakhun and their English translations:\n"
|
|
"1. \u014ba ka k\u0264 ne | Do I go?\n"
|
|
"2. n\u0264 \u02b7ip tu\u0294 ne | Did you sleep?\n"
|
|
"3. \u014bab\u0259 ati lapk\u02b0i t\u0264\u0294 ne | Did I see him?\n"
|
|
"10. ati k\u0259m\u0259 \u014ba lapk\u02b0i t\u02b0\u0264 ne | Did he see me?"
|
|
),
|
|
"query": "Translate into English:\n1. n\u0264 \u02b7ip ku ne\n2. ati k\u0259m\u0259 nirum lapk\u02b0i t\u02b0i ne",
|
|
"answer": "Do you sleep?\nDid he see us?",
|
|
},
|
|
"match_letters": {
|
|
"context": (
|
|
"Given are words in Nahuatl as well as their English translations in arbitrary order:\n"
|
|
"1. acalhuah A. water\n2. achilli B. child\n3. atl C. master of house\n"
|
|
"4. callah D. water pepper\n18. totoltetl R. revered grandfather"
|
|
),
|
|
"query": "Determine the correct correspondences.",
|
|
"answer": "O\nD\nA\nG\nC\nH",
|
|
},
|
|
"fill_blanks": {
|
|
"context": (
|
|
"Here are two different forms of some verbs in Guazacap\u00e1n Xinka and their English translations:\n"
|
|
"piriy\u02bc | \u0268mbir\u02bci | see\nim\u02bcay | \u0268nim\u02bca | say, tell\n"
|
|
"k\u02bcaniy\u02bc | \u0268\u014bk\u02bcan\u02bci | trap\nter\u02bcoy | \u0268nder\u02bco | kill"
|
|
),
|
|
"query": "Fill the blanks (1-2):\nnetkay\u02bc | (1) | push\nk\u0268r\u0268y\u02bc | (2) | pull",
|
|
"answer": "\u0268nnetak\u02bca\n\u0268\u014bgir\u02bci",
|
|
},
|
|
"text_to_num": {
|
|
"context": (
|
|
"The squares of the numbers 1 to 10 are spelt out in the Ndom language, in arbitrary order:\n"
|
|
"nif abo mer an thef abo sas\nnif thef abo tondor abo mer abo thonith\n"
|
|
"mer an thef abo thonith\nmer abo ithin"
|
|
),
|
|
"query": "Write in numerals:\n1. nif ithin abo ithin\n2. mer an thef abo meregh",
|
|
"answer": "111\n17",
|
|
},
|
|
}
|
|
FEWSHOT_EXAMPLES["num_to_text"] = FEWSHOT_EXAMPLES["text_to_num"]
|
|
|
|
|
|
def build_messages(system, context, query, task_type=None):
|
|
messages = [{"role": "system", "content": system}]
|
|
ex = FEWSHOT_EXAMPLES.get(task_type) if task_type else None
|
|
if ex:
|
|
messages.append({"role": "user", "content": f"{ex['context']}\n\n{ex['query']}"})
|
|
messages.append({"role": "assistant",
|
|
"content": f"(reasoning omitted for brevity)\n\nFINAL ANSWERS:\n{ex['answer']}"})
|
|
messages.append({"role": "user", "content": f"{context.strip()}\n\n{query.strip()}"})
|
|
return messages
|
|
|
|
|
|
def count_items(query: str) -> int:
|
|
nums = re.findall(r"(?m)^\s*\(?(\d+)[.)]\s", query)
|
|
return len(set(nums)) if nums else 1
|
|
|
|
|
|
def parse_answers(text: str, n_expected: int):
|
|
marker = list(re.finditer(r"(?im)^\s*final answers?\s*:?\s*$", text))
|
|
tail = text[marker[-1].end():] if marker else text
|
|
answers = []
|
|
for line in tail.splitlines():
|
|
line = re.sub(r"^\s*\(?\d+[.)]\s*", "", line).strip().strip('"').strip("'")
|
|
if line:
|
|
answers.append(line)
|
|
if len(answers) < n_expected:
|
|
answers += [""] * (n_expected - len(answers))
|
|
elif len(answers) > n_expected:
|
|
answers = answers[:n_expected]
|
|
return answers
|
|
|
|
|
|
def vote(candidate_lists, n_expected):
|
|
out = []
|
|
for pos in range(n_expected):
|
|
counts = Counter()
|
|
display = {}
|
|
for cand in candidate_lists:
|
|
v = cand[pos] if pos < len(cand) else ""
|
|
key = v.strip().lower()
|
|
if key:
|
|
counts[key] += 1
|
|
display.setdefault(key, v)
|
|
if counts:
|
|
best_key = counts.most_common(1)[0][0]
|
|
out.append(display[best_key])
|
|
else:
|
|
out.append(candidate_lists[0][pos] if pos < len(candidate_lists[0]) else "")
|
|
return out
|
|
|
|
|
|
def generate(model, tok, messages, max_new_tokens, do_sample=False, temperature=0.7):
|
|
enc = tok.apply_chat_template(
|
|
messages, add_generation_prompt=True, return_tensors="pt", return_dict=True,
|
|
).to(model.device)
|
|
gen_kwargs = dict(max_new_tokens=max_new_tokens)
|
|
if do_sample:
|
|
gen_kwargs.update(do_sample=True, temperature=temperature, top_p=0.9)
|
|
else:
|
|
gen_kwargs.update(do_sample=False)
|
|
with torch.no_grad():
|
|
out = model.generate(**enc, **gen_kwargs)
|
|
return tok.decode(out[0][enc["input_ids"].shape[-1]:], skip_special_tokens=True).strip()
|
|
|
|
|
|
def main():
|
|
start = time.time()
|
|
tok = AutoTokenizer.from_pretrained(MODEL_ID)
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
MODEL_ID, torch_dtype=torch.float16, device_map="auto",
|
|
).eval()
|
|
|
|
rows = list(csv.DictReader(open("/tmp/data/test.csv", newline="")))
|
|
out_rows = []
|
|
|
|
for idx, r in enumerate(rows):
|
|
context = (r.get("context") or "").strip()
|
|
query = (r.get("query") or "").strip()
|
|
task_type = r.get("task_type")
|
|
n_exp = count_items(query)
|
|
messages = build_messages(SYSTEM, context, query, task_type)
|
|
|
|
# Time-budget throttling: once we're low on time, fall back to a single
|
|
# greedy pass with a shorter budget instead of risking a hard timeout.
|
|
elapsed = time.time() - start
|
|
remaining_rows = max(1, len(rows) - idx)
|
|
time_left = TIME_LIMIT_SECONDS - elapsed
|
|
per_row_budget = time_left / remaining_rows if remaining_rows else 0
|
|
tight = per_row_budget < 20 or time_left < 60
|
|
k_samples = 1 if tight else K_SAMPLES
|
|
max_new = 384 if tight else MAX_NEW_TOKENS
|
|
|
|
try:
|
|
raws, candidates = [], []
|
|
for k in range(k_samples):
|
|
raw = generate(model, tok, messages, max_new, do_sample=(k > 0), temperature=TEMPERATURE)
|
|
raws.append(raw)
|
|
candidates.append(parse_answers(raw, n_exp))
|
|
preds = vote(candidates, n_exp)
|
|
raw0 = raws[0]
|
|
except Exception as e:
|
|
preds = [""] * n_exp
|
|
raw0 = f"[error: {e}]"
|
|
|
|
row_out = {"id": r["id"], "pred": json.dumps(preds, ensure_ascii=False)}
|
|
|
|
if WRITE_EXPLANATIONS and (time.time() - start) < TIME_LIMIT_SECONDS - 30:
|
|
try:
|
|
explanation = generate(model, tok,
|
|
[{"role": "system", "content": SUMMARIZE},
|
|
{"role": "user", "content": raw0}],
|
|
200)
|
|
except Exception:
|
|
explanation = ""
|
|
row_out["explanation"] = explanation
|
|
elif WRITE_EXPLANATIONS:
|
|
row_out["explanation"] = ""
|
|
|
|
out_rows.append(row_out)
|
|
print(f"{idx + 1}/{len(rows)} done (k={k_samples}, elapsed={elapsed:.0f}s)", flush=True)
|
|
|
|
fieldnames = ["id", "pred", "explanation"] if WRITE_EXPLANATIONS else ["id", "pred"]
|
|
with open("submission.csv", "w", newline="") as f:
|
|
w = csv.DictWriter(f, fieldnames=fieldnames)
|
|
w.writeheader()
|
|
for row in out_rows:
|
|
w.writerow(row)
|
|
print("wrote submission.csv", flush=True)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|