import os import subprocess import sys def _install_bundled_deps() -> None: """Install transformers from bundled wheels (eval sandbox has no PyPI access).""" wheels_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "wheels") if not os.path.isdir(wheels_dir): return subprocess.run( [ sys.executable, "-m", "pip", "install", "-q", "--no-index", f"--find-links={wheels_dir}", "transformers==4.56.2", ], check=True, ) _install_bundled_deps() import re import csv import json import shutil import tempfile import torch from transformers import AutoTokenizer, AutoModelForCausalLM # 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 = "." MAX_NEW_TOKENS = 2000 TEMPERATURE = 0.8 TOP_P = 0.95 MAX_ATTEMPTS = 5 def load_tokenizer(model_id: str = "."): """Load tokenizer, converting tokenizer.json for older tokenizers if needed.""" tokenizer_path = os.path.join(model_id, "tokenizer.json") with open(tokenizer_path, encoding="utf-8") as handle: data = json.load(handle) merges = data.get("model", {}).get("merges", []) if not merges or not isinstance(merges[0], list): return AutoTokenizer.from_pretrained(model_id) # Older tokenizers expect merge pairs as "a b" strings, not ["a", "b"] lists. data["model"]["merges"] = [" ".join(piece) for piece in merges] tmpdir = tempfile.mkdtemp() for name in ("tokenizer_config.json", "special_tokens_map.json"): src = os.path.join(model_id, name) if os.path.isfile(src): shutil.copy(src, tmpdir) with open(os.path.join(tmpdir, "tokenizer.json"), "w", encoding="utf-8") as handle: json.dump(data, handle) return AutoTokenizer.from_pretrained(tmpdir) 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.\n" "Structure your response in exactly two sections with these headers:\n\n" "**Reasoning:**\n" "(your step-by-step analysis of linguistic rules, how to apply them, and answer format)\n\n" "**Final Answers:**\n" "(one answer per line, in query order -- bare answers only, no numbering, no quotes, " "no extra text, according to the TASK TYPE)" ) tok = load_tokenizer(MODEL_ID) model = AutoModelForCausalLM.from_pretrained( MODEL_ID, torch_dtype=torch.float16, device_map="auto" ).eval() with open("/tmp/data/test.csv", encoding="utf-8", newline="") as f: test_rows = list(csv.DictReader(f)) outputs_queries_types = [] for r in test_rows: # Create the prompt. messages = [ {"role": "system", "content": SYSTEM}, {"role": "user", "content": f"CONTEXT:{r['context'].strip()}\nTASK TYPE:`{r['task_type']}`\n\nQUERY:{r['query'].strip()}"}, ] ids = tok.apply_chat_template( messages, add_generation_prompt=True, return_tensors="pt", ).to(model.device) # Generate the answer. done = False attempts = 0 while not done and attempts < MAX_ATTEMPTS: with torch.no_grad(): out = model.generate( ids, max_new_tokens=MAX_NEW_TOKENS, do_sample=True, temperature=TEMPERATURE, top_p=TOP_P,) text = tok.decode(out[0][ids.shape[-1]:], skip_special_tokens=True).strip() # IF NO FINAL ANSWER keyword is used, try again. attempts += 1 if "final answer" not in text.lower() or text.lower().split('final answer')[1].split('\n')==0: print(f'TRYING AGAIN...attempts #{attempts+1}/{MAX_ATTEMPTS}') else: done = True outputs_queries_types.append((text, r['id'], r['query'], r['task_type'])) print(f"{len(outputs_queries_types)}/{len(test_rows)} done", flush=True) # Postprocess and store the answers. # Generous line-start markers (same spirit as base script): allow markdown, # headers, etc. REASONING_MARKER_RE = re.compile( r"(?im)^[^\w\n]*(?:step-by-step\s+)?reasoning[^\w\n]*:?\s*$" ) FINAL_ANSWERS_MARKER_RE = re.compile( r"(?im)^[^\w\n]*final answers?[^\w\n]*:?\s*$" ) def split_response(text: str) -> tuple[str, str]: """Return (explanation, raw final-answers section text).""" reasoning_matches = list(REASONING_MARKER_RE.finditer(text)) final_matches = list(FINAL_ANSWERS_MARKER_RE.finditer(text)) raw_final = "" if final_matches: raw_final = text[final_matches[-1].end() :].strip() explanation = "" if reasoning_matches and final_matches: reasoning_start = reasoning_matches[0].end() finals_after_reasoning = [ match for match in final_matches if match.start() >= reasoning_start ] if finals_after_reasoning: explanation = text[reasoning_start : finals_after_reasoning[0].start()].strip() if not explanation: explanation = raw_final return explanation, raw_final def expected_answer_count(query: str, task_type: str) -> int: if task_type == "match_letters": numbered = re.findall(r"^\s*\d+\.", query, re.MULTILINE) return len(numbered) or 1 if "blanks" in query.lower(): range_match = re.search(r"\((\d+)-(\d+)\)", query) if range_match: return int(range_match.group(2)) - int(range_match.group(1)) + 1 return len(re.findall(r"\(\d+\)", query)) or 1 numbered = re.findall(r"^\s*\d+[.)]", query, re.MULTILINE) return len(numbered) or 1 def split_single_line_answer(text: str, expected: int, task_type: str) -> list[str]: text = text.strip() if expected <= 1: return [text] def try_split(pattern: str) -> list[str] | None: parts = [part.strip() for part in re.split(pattern, text) if part.strip()] return parts if len(parts) == expected else None if task_type == "match_letters": for pattern in (r"\s+", r",\s*", r";\s*"): if result := try_split(pattern): return result letters = re.findall(r"[A-Za-z]", text) if len(letters) == expected: return [letter.upper() for letter in letters] return [text] if task_type in ("text_to_num", "num_to_text"): for pattern in (r",\s*", r";\s*", r"\s+"): if result := try_split(pattern): return result return [text] for pattern in (r";\s*", r",\s*"): if result := try_split(pattern): return result return [text] def parse_answer_lines(text_after_marker: str, query: str, task_type: str) -> list[str]: """Parse cleaned answer lines from the raw final-answers section.""" answers = [] for line in text_after_marker.splitlines(): stripped_line = line.strip("`").strip() if stripped_line == "": continue match_numbered_prefix = re.match(r"^\s*\d+[.)]\s+(.*)", stripped_line) if match_numbered_prefix: cleaned_line = match_numbered_prefix.group(1).strip() else: cleaned_line = stripped_line cleaned_line = re.sub(r"\*\*", "", cleaned_line).strip() if task_type == "match_letters": parts = [ part.strip("().[]") for part in re.split(r"[\s,;]+", cleaned_line) if part.strip() ] if not ( len(parts) > 1 and all(re.fullmatch(r"[A-Za-z]", part) for part in parts) ): match_letter_word = re.match( r"^\s*(?:\(([A-Za-z])\)|\[([A-Za-z])\]|([A-Za-z]))\.?:?\s*(.*)$", cleaned_line, ) if match_letter_word: letter = ( match_letter_word.group(1) or match_letter_word.group(2) or match_letter_word.group(3) ) cleaned_line = letter.upper() if cleaned_line: answers.append(cleaned_line) expected = expected_answer_count(query, task_type) if len(answers) == 1 and expected > 1: answers = split_single_line_answer(answers[0], expected, task_type) return answers def postprocess_answer(text, query, task_type): """Extract explanation and parsed final answers from model output.""" explanation, raw_final = split_response(text) if not raw_final: return [], explanation return parse_answer_lines(raw_final, query, task_type), explanation rows = [] for answer, row_id, query, task_type in outputs_queries_types: answers, explanation = postprocess_answer(answer, query, task_type) rows.append( { "id": row_id, "pred": json.dumps(answers, ensure_ascii=False), "explanation": explanation, } ) with open("submission.csv", "w", encoding="utf-8", newline="") as f: writer = csv.DictWriter(f, fieldnames=["id", "pred", "explanation"]) writer.writeheader() writer.writerows(rows) print("wrote submission.csv", flush=True)