#import subprocess, sys #subprocess.run( # [ # sys.executable, # "-m", # "pip", # "install", # "-q", # "transformers>=4.43", # "accelerate>=0.30", # "torch>=2.2", # "pandas", # ], # check=True, #) import os import re import json import pandas as pd import torch from transformers import AutoTokenizer, AutoModelForCausalLM from incontext_examples import INCONTEXT_EXAMPLES # The repository is the working directory at runtime. # The evaluation sandbox has no internet access. os.environ["HF_HUB_OFFLINE"] = "1" os.environ["TRANSFORMERS_OFFLINE"] = "1" MODEL_ID = "." TEST_PATH = "/tmp/data/test.csv" OUTPUT_PATH = "submission.csv" MAX_NEW_TOKENS = 1536 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID) if tokenizer.pad_token_id is None: tokenizer.pad_token = tokenizer.eos_token model = AutoModelForCausalLM.from_pretrained( MODEL_ID, torch_dtype=torch.float16, device_map="auto", ).eval() df = pd.read_csv( TEST_PATH, dtype=str, ).fillna("") required_columns = { "id", "context", "query", "task_type", } missing_columns = required_columns - set(df.columns) if missing_columns: raise ValueError( "The test CSV is missing required columns: " f"{sorted(missing_columns)}" ) SYSTEM = ( "You solve International Linguistics Olympiad problems by reasoning from the " "data you are given. You may meet a task type you have never seen: read the " "instruction and the examples, and answer in the same form they use. " "Common task types and what to give -- " "translation: the translated form only, in the language the task asks for; " "fill_blanks: only the missing form for each blank; " "match_letters: only the option letter, for example A, B, or C; " "text_to_num: the number in digits; " "num_to_text: the number written out in words, in the language asked; " "any other type: give exactly what the instruction asks, nothing else. " "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. " "Each final answer must be bare: no numbering, no quotes, and no extra text." ) def normalize_task_type(task_type): """ Normalize task-type names so capitalization, spaces, and hyphens do not prevent a match. Examples: "Num To Text" -> "num_to_text" "fill-blanks" -> "fill_blanks" """ task_type = str(task_type).strip().lower() task_type = re.sub(r"[\s-]+", "_", task_type) return task_type def parse_answers(text): """ Keep only nonempty lines after the final FINAL ANSWERS marker. If the model adds numbering such as: 1. answer 2) answer the numbering is removed. """ markers = list( re.finditer( r"(?im)^\s*final answers?\s*:?\s*$", text, ) ) if markers: text = text[markers[-1].end():] answers = [] for line in text.splitlines(): line = re.sub( r"^\s*\d+[.)]\s*", "", line, ).strip() if line: answers.append(line) return answers def format_example_assistant_response(example): """ Format the stored reasoning trace and answers as an assistant message. """ reasoning = str( example.get("reasoning", "") ).strip() answers = example.get("answers", []) if isinstance(answers, str): answers = [ line.strip() for line in answers.splitlines() if line.strip() ] formatted_answers = "\n".join( str(answer).strip() for answer in answers ) if reasoning: return ( f"{reasoning}\n\n" f"FINAL ANSWERS:\n" f"{formatted_answers}" ) return ( f"FINAL ANSWERS:\n" f"{formatted_answers}" ) def get_incontext_example(task_type): """ Retrieve the demonstration matching the current problem's task type. Returns None when no demonstration exists for that task type. """ normalized_type = normalize_task_type(task_type) # Normalize the dictionary keys as well, so keys such as # "num-to-text" and "num_to_text" are treated equivalently. normalized_examples = { normalize_task_type(key): value for key, value in INCONTEXT_EXAMPLES.items() } return normalized_examples.get(normalized_type) def build_messages(row): """ Construct the conversation for one test problem. The conversation contains: 1. the system prompt; 2. a solved demonstration of the same task type, when available; 3. the actual test problem. """ task_type = normalize_task_type( row["task_type"] ) example = get_incontext_example( task_type ) messages = [ { "role": "system", "content": SYSTEM, } ] if example is not None: example_context = str( example.get("context", "") ).strip() example_query = str( example.get("query", "") ).strip() messages.extend( [ { "role": "user", "content": ( f"Task type: {task_type}\n\n" f"{example_context}\n\n" f"{example_query}" ), }, { "role": "assistant", "content": format_example_assistant_response( example ), }, ] ) else: print( "Warning: no in-context example found for " f"task_type={task_type!r}, id={row['id']!r}. " "Running the problem without a demonstration." ) problem_context = str( row["context"] ).strip() problem_query = str( row["query"] ).strip() messages.append( { "role": "user", "content": ( f"Task type: {task_type}\n\n" f"{problem_context}\n\n" f"{problem_query}" ), } ) return messages rows = [] for _, row in df.iterrows(): messages = build_messages( row ) inputs = tokenizer.apply_chat_template( messages, add_generation_prompt=True, return_tensors="pt", return_dict=True, ).to(model.device) input_length = inputs[ "input_ids" ].shape[-1] with torch.no_grad(): output = model.generate( **inputs, max_new_tokens=MAX_NEW_TOKENS, do_sample=False, pad_token_id=tokenizer.pad_token_id, eos_token_id=tokenizer.eos_token_id, ) generated_text = tokenizer.decode( output[0][input_length:], skip_special_tokens=True, ).strip() answers = parse_answers( generated_text ) rows.append( { "id": row["id"], "pred": json.dumps( answers, ensure_ascii=False, ), } ) submission_df = pd.DataFrame( rows ) submission_df.to_csv( OUTPUT_PATH, index=False, ) print( f"Saved {len(submission_df)} predictions to {OUTPUT_PATH}" )