| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
|
|
| import os |
| import re |
| import json |
|
|
| import pandas as pd |
| import torch |
| from transformers import AutoTokenizer, AutoModelForCausalLM |
|
|
| from incontext_examples import INCONTEXT_EXAMPLES |
|
|
|
|
| |
| |
| 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) |
|
|
| |
| |
| 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}" |
| ) |
