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)