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()