""" Competition submission script — copy this into your HF repo as script.py. The eval sandbox: - mounts the test set at /tmp/data/test.csv - has no internet - runs on a T4 (16GB) - has 30 minutes - has bitsandbytes and autoawq pre-installed Strategy: Ship Qwen2.5-14B-Instruct-AWQ with HYBRID prompting. The 14B-AWQ is the proven competition baseline (0.123 score). We improve on the baseline with: 1. Task-specific CoT prompts for translation/fill_blanks (improves EM) 2. Direct prompts for match_letters/text_to_num/num_to_text (faster) 3. Adaptive max_new_tokens per task type (512 for CoT, 256 for direct) 4. Fixed answer parser (v1 dropped ~5% of correct answers) 5. Explanation column for human jury track 6. Time guard to never exceed 30-min limit Why hybrid: Pure CoT was too slow (70s/problem) and exceeded the time budget. Pure direct prompting gave EM=0.025 on the hidden test set. CoT for hard tasks (translation, fill_blanks) improves exact matches by letting the model reason carefully; direct prompting is fine for pattern-matching tasks where reasoning doesn't help. """ import json import os import re import time os.environ["HF_HUB_OFFLINE"] = "1" os.environ["TRANSFORMERS_OFFLINE"] = "1" MODEL_ID = "." import pandas as pd import torch from transformers import AutoTokenizer, AutoModelForCausalLM from prompts import ( get_system_prompt, USER_TEMPLATE, parse_answers, extract_analysis, count_query_items, ) # Time budget: 30 min total. Reserve 3 min for model loading + CSV write. TIME_BUDGET_S = 27 * 60 # 27 minutes for inference def load_model(): """Load the AWQ-quantized model for T4 16GB.""" tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True) if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token model = AutoModelForCausalLM.from_pretrained( MODEL_ID, device_map="auto", trust_remote_code=True, torch_dtype=torch.float16, ) model.eval() print("[submit] Loaded Qwen2.5-14B-Instruct-AWQ", flush=True) return tokenizer, model def solve_problem( tokenizer, model, context: str, query: str, task_type: str = "", max_new_tokens: int = 256, ) -> tuple[list[str], str]: """Generate answers for one IOL problem. For translation/fill_blanks: CoT reasoning (max 512 tokens). For match_letters/text_to_num/num_to_text: direct (256 tokens). """ n_items = count_query_items(query) system_prompt = get_system_prompt(task_type) messages = [ {"role": "system", "content": system_prompt}, {"role": "user", "content": USER_TEMPLATE.format( context=context.strip(), query=query.strip() )}, ] text = tokenizer.apply_chat_template( messages, add_generation_prompt=True, tokenize=False ) inputs = tokenizer(text, return_tensors="pt") input_ids = inputs["input_ids"].to(model.device) with torch.no_grad(): # Greedy decoding (do_sample=False) — reproducible and best for our use case. # Beam search (num_beams=2) tested but caused catastrophic failures (0.0 score). out = model.generate( input_ids, max_new_tokens=max_new_tokens, do_sample=False, pad_token_id=tokenizer.eos_token_id, ) generated = tokenizer.decode( out[0][input_ids.shape[-1]:], skip_special_tokens=True ).strip() answers = parse_answers(generated, n_expected=n_items, task_type=task_type) explanation = extract_analysis(generated) return answers, explanation def _format_pred(answers: list[str]) -> str: """Format predictions for submission. Output is JSON-encoded list of answer strings (the IOL competition evaluator parses this with ast.literal_eval). We also include a pipe-separated fallback in a comment-like column for safety. """ return json.dumps(answers, ensure_ascii=False) def main(): t_start = time.time() print("[submit] Loading model...", flush=True) tokenizer, model = load_model() print("[submit] Reading test set...", flush=True) df = pd.read_csv("/tmp/data/test.csv", dtype=str).fillna("") n_problems = len(df) print(f"[submit] Loaded {n_problems} problems", flush=True) # Estimate time per problem type for adaptive budget # CoT tasks: ~20s each (512 tokens); direct tasks: ~5s each (256 tokens) COT_TASKS = {"translation", "fill_blanks"} DIRECT_TASKS = {"match_letters", "text_to_num", "num_to_text"} SHORT_TASKS = {"match_letters", "text_to_num"} # single char / digits rows = [] for idx, row in df.iterrows(): elapsed = time.time() - t_start remaining = TIME_BUDGET_S - elapsed problems_left = n_problems - idx task_type = row.get("task_type", "") # Adaptive max_new_tokens based on time remaining and task type # Target: average ~10s per problem to fit 160 problems in 27 min # Budget per problem: 27*60 / 160 = 10.1s # CoT is ~17s, direct is ~5-8s, short is ~4s # Be aggressive: switch to fallback when remaining < 8s/problem if remaining < problems_left * 8 and remaining > 0 and idx > 0: # Tight on time — minimal tokens, direct mode (no CoT) current_max = 96 use_cot = False if idx % 10 == 0: print(f"[submit] FAST MODE at {idx+1}/{n_problems} " f"({remaining:.0f}s left, {remaining/problems_left:.1f}s/problem)", flush=True) elif remaining < problems_left * 12 and remaining > 0 and idx > 0 and task_type in COT_TASKS: # Getting tight on CoT problems — reduce CoT max tokens current_max = 256 use_cot = True if idx % 10 == 0: print(f"[submit] COOL DOWN at {idx+1}/{n_problems} " f"({remaining:.0f}s left, {remaining/problems_left:.1f}s/problem)", flush=True) elif task_type in COT_TASKS: # Verbose CoT (matches 0.0872 baseline that scored highest). # 512 tokens gives room for full step-by-step reasoning. current_max = 512 use_cot = True elif task_type in SHORT_TASKS: # Short answers (single letters or digits) — keep tight current_max = 128 use_cot = False else: # Direct for easy tasks (num_to_text, etc.) current_max = 256 use_cot = False try: # If we need to force direct mode for time, swap to default prompt if not use_cot and task_type in COT_TASKS and remaining < problems_left * 8: # Switch to default prompt (direct) for time-constrained CoT tasks from prompts import _DEFAULT_PROMPT original_prompt = get_system_prompt(task_type) # Use default prompt via monkey-patch import prompts prompts._PROMPTS[task_type] = _DEFAULT_PROMPT answers, explanation = solve_problem( tokenizer, model, context=row["context"], query=row["query"], task_type=task_type, max_new_tokens=current_max, ) prompts._PROMPTS[task_type] = original_prompt else: answers, explanation = solve_problem( tokenizer, model, context=row["context"], query=row["query"], task_type=task_type, max_new_tokens=current_max, ) except Exception as e: print(f"[submit] ERROR at {idx+1}/{n_problems}: {e}", flush=True) n_items = count_query_items(row.get("query", "")) answers = [""] * max(n_items, 1) explanation = "" rows.append({ "id": row["id"], "pred": _format_pred(answers), "explanation": explanation, }) if (idx + 1) % 10 == 0 or idx == 0: print(f"[submit] {idx + 1}/{n_problems} done " f"({elapsed:.0f}s elapsed, task={task_type})", flush=True) output = pd.DataFrame(rows) # Write to the path the eval system expects import os as _os _os.makedirs("/tmp/model", exist_ok=True) output.to_csv("/tmp/model/submission.csv", index=False) # Also write a backup at the relative path (in case CWD is /tmp/model) output.to_csv("submission.csv", index=False) total_elapsed = time.time() - t_start print(f"[submit] wrote submission.csv ({len(rows)} problems, " f"{total_elapsed:.0f}s total, cwd={_os.getcwd()})", flush=True) if __name__ == "__main__": main()