| 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 |
|
|
| |
| 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) |
|
|
| |
| 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." |
| "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." |
| ) |
|
|
| SYSTEM_POST_EXPLAIN = ( |
| "You are a helpful assistant that explains the reasoning behind the answers to the International Linguistics Olympiad problems. |
| "You are given the following information:\n" |
| "- The context of the problem\n" |
| "- The task type\n" |
| "- The query\n" |
| "- The answer\n" |
| "- The reasoning\n" |
| "You need to explain the reasoning behind the answer in a way that is easy to understand and concisely focused on the key insights and rules deduced and applied.\n" |
| "Do not include any other text, do not includethe answer in the explanation, " |
| " and do not invent any new information." |
| ) |
| 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 |
| |
| # Generate the explanation. |
| messages_post_explain = [ |
| {"role": "system", "content": SYSTEM_POST_EXPLAIN}, |
| {"role": "user", "content": |
| f"CONTEXT:{r['context'].strip()}\nTASK TYPE:`{r['task_type']}`\n\nQUERY:{r['query'].strip()}\n\nANSWER:{text.strip()}"}, |
| ] |
| ids_post_explain = tok.apply_chat_template( |
| messages_post_explain, add_generation_prompt=True, return_tensors="pt", |
| ).to(model.device) |
| with torch.no_grad(): |
| out_post_explain = model.generate( |
| ids_post_explain, |
| max_new_tokens=MAX_NEW_TOKENS, |
| do_sample=False) |
| text_post_explain = tok.decode(out_post_explain[0][ids_post_explain.shape[-1]:], skip_special_tokens=True).strip() |
| |
| outputs_queries_types.append((text, text_post_explain, r['id'], r['query'], r['task_type'])) |
| print(f"{len(outputs_queries_types)}/{len(test_rows)} done", flush=True) |
| |
| # Postprocess and store the answers. |
| 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 postprocess_answer(text, query, task_type): |
| """Keep only the lines after the last 'FINAL ANSWERS:' marker, one answer per line, |
| stopping at the first empty line. If eval_type is multiple and only one line as answer, split at whitespace.""" |
| # Updated regex to be more flexible with surrounding characters |
| marker_match = list(re.finditer(r"(?im)^[^\w\n]*final answers?[^\w\n]*:?\s*$", text)) |
| if marker_match: |
| text_after_marker = text[marker_match[-1].end():] |
| #print('FOUND FINAL ANSWER', text_after_marker) |
| else: |
| #print("No 'FINAL ANSWERS:' marker found") |
| return [] |
| |
| answers = [] |
| answer_lines = text_after_marker.splitlines() |
| for i, line in enumerate(answer_lines): |
| stripped_line = line.strip('`').strip() |
| |
| # Stop processing if an empty line is encountered (not as first line) |
| if stripped_line=='': |
| continue |
| |
| # Use a more precise regex to only remove numbering if it's a prefix to other text |
| # This ensures that lines which are just numbers (e.g., '1') are not stripped. |
| 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 |
| |
| # Remove any bold markdown '**' |
| cleaned_line = re.sub(r"\*\*", "", cleaned_line).strip() |
| |
| # Specific handling for 'match_letters' task type to strip extra words |
| 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() |
| |
| # Append the cleaned, non-empty line |
| if cleaned_line: |
| answers.append(cleaned_line) |
| |
| #print('PARSED ANSWERS', answers) |
| |
| # Compare against QUERY length: sometimes model forgets newlines |
| #print('QUERY', query) |
| expected = expected_answer_count(query, task_type) |
| query_len = len(query.splitlines()) - 2 |
| #print(query_len) |
| if len(answers) == 1 and expected > 1: |
| answers = split_single_line_answer(answers[0], expected, task_type) |
| return answers |
| |
| rows = [] |
| for answer, explanation, row_id, query, task_type in outputs_queries_types: |
| answers = 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) |
