import os os.environ["HF_HUB_OFFLINE"] = "1" os.environ["TRANSFORMERS_OFFLINE"] = "1" # --- make sure quantization lib is present (PyPI is reachable at run time) --- import subprocess, sys def _run_pkg_cmd(args): for cmd in (["uv", "pip", *args], [sys.executable, "-m", "pip", *args], ["pip", *args]): try: if subprocess.run(cmd, capture_output=True, text=True).returncode == 0: return True except FileNotFoundError: continue return False try: import bitsandbytes # noqa: F401 except ImportError: _run_pkg_cmd(["install", "-q", "bitsandbytes"]) # NOTE: torchvision is uninstalled here because it can pull in a conflicting # pinned torch version when transformers/bitsandbytes resolve dependencies. # We don't use any vision functionality, so this is safe. try: import torchvision # noqa: F401 _run_pkg_cmd(["uninstall", "-y", "-q", "torchvision"]) except ImportError: pass import re import json import time import math import os import sys import pandas as pd import torch from collections import Counter from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig # Safety configuration DISABLE_RL = os.environ.get("IOL_SAFE_MODE", "").lower() == "true" if DISABLE_RL: print("SAFE MODE: RL and HP tuning disabled", flush=True) MODEL_ID = "." TIME_LIMIT = 30 * 60 # hard competition limit, seconds SAFETY_BUFFER = 90 # stop issuing new generations this many seconds before the limit MAX_NEW_TOKENS = 1800 # increased for scratchpad + rules NUM_PATHS = 3 # N=3 for test-time scaling TEMPERATURE = 0.3 # T=0.3 for diverse but coherent paths TOP_P = 0.9 # default nucleus sampling value (FIX: was only ever set as a # side-effect global inside apply_tuned_hyperparams and never # actually read by generate_n_paths) START = time.time() def time_left(): return TIME_LIMIT - (time.time() - START) print("Loading tokenizer/model...", flush=True) tok = AutoTokenizer.from_pretrained(MODEL_ID) bnb_config = BitsAndBytesConfig( load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.float16, bnb_4bit_use_double_quant=True, ) # Load model with OOM protection try: model = AutoModelForCausalLM.from_pretrained( MODEL_ID, device_map="auto", torch_dtype=torch.float32, ).eval() print("Model loaded successfully", flush=True) except torch.cuda.OutOfMemoryError as e: print(f"FATAL: GPU OOM during model load: {e}", flush=True) # Try CPU fallback try: print("Attempting CPU loading...", flush=True) model = AutoModelForCausalLM.from_pretrained( MODEL_ID, device_map="auto", torch_dtype=torch.float32, ).eval() except Exception as e2: print(f"FATAL: Could not load model even on CPU: {e2}", flush=True) # Create error submission with open("submission.csv", "w") as f: f.write("id,pred,explanation\n") sys.exit(1) print(f"Model loaded in {time.time() - START:.1f}s", flush=True) # Robust data loading with error handling try: if not os.path.exists("/tmp/data/test.csv"): print("ERROR: test.csv not found at /tmp/data/test.csv", flush=True) # Create empty submission as fallback with open("submission.csv", "w") as f: f.write("id,pred,explanation\n") sys.exit(1) df = pd.read_csv("/tmp/data/test.csv", dtype=str).fillna("") if len(df) == 0: print("WARNING: test.csv is empty", flush=True) # Still create valid submission pd.DataFrame(columns=["id", "pred", "explanation"]).to_csv("submission.csv", index=False) sys.exit(0) print(f"Loaded {len(df)} test items", flush=True) except Exception as e: print(f"FATAL: Could not load test.csv: {e}", flush=True) # Create error submission with open("submission.csv", "w") as f: f.write("id,pred,explanation\n") sys.exit(1) ITEM_RE = re.compile(r"(?m)^\s*(\d+)\.\s") # ============================================================================= # 1. DYNAMIC TASK ROUTER # Analyzes: task_type, eval_type, item count (K) # ============================================================================= TASK_PROFILES = { "translation": { "eval_metric": "chrF", "format_hint": "Translate to the target language preserving diacritics and orthography.", "grammar_focus": "morphophonology, tense, case, agreement", }, "match_letters": { "eval_metric": "exact", "format_hint": "Answer with only the option letter (e.g. A, B, C) for each item.", "grammar_focus": "phonological rules, allophony, orthographic patterns", }, "fill_blanks": { "eval_metric": "chrF", "format_hint": "Fill the blank with the correct inflected/corrected form.", "grammar_focus": "inflection, derivation, agreement, sandhi", }, "text_to_num": { "eval_metric": "exact", "format_hint": "Answer with the value written in digits (e.g. 285).", "grammar_focus": "numeral systems, base systems, place value", }, "num_to_text": { "eval_metric": "chrF", "format_hint": "Answer with the number written out in words in the task language.", "grammar_focus": "numeral morphology, ordinals, cardinals", }, } def count_items(query: str) -> int: """Extract K = number of items in the query. FIX: previously used len(set(nums)) which under-counts whenever a numeric label repeats or numbering is non-contiguous (e.g. duplicated "1." lines, grouped sub-items). Using the max item index found is robust to that, since IOL-style queries number items sequentially starting at 1. """ nums = ITEM_RE.findall(query) if nums: return max(int(n) for n in nums) # fallback: count non-empty lines return max(1, len([l for l in query.splitlines() if l.strip()])) def get_task_profile(task_type: str) -> dict: """Route to task-specific configuration.""" return TASK_PROFILES.get(task_type, TASK_PROFILES["translation"]) # ============================================================================= # 2. FEW-SHOT EXAMPLES: Filled-out Grammar Scratchpads # Shows the model HOW to reason, not just WHAT to do # ============================================================================= FEW_SHOT_EXAMPLES = """ === EXAMPLE 1: Simple Translation === DATA: 1. mo ka | I go 2. ti ka | you go 3. mo tuʔ | I sleep 4. ti tuʔ | you sleep QUERY: Translate into English: 1. ti tuʔ 2. mo ka 【GRAMMAR SCRATCHPAD】 1. DATA INVENTORY: - mo: "I" (appears in sentences 1, 3) - ti: "you" (appears in sentences 2, 4) - ka: "go" (appears in sentences 1, 2) - tuʔ: "sleep" (appears in sentences 3, 4) 2. MORPHEME SEGMENTATION: - All words are monomorphemic (no affixes detected) - Structure: [PRONOUN] + [VERB] 3. PARADIGM MAPPING: - Pronouns: mo = 1st person singular, ti = 2nd person singular - Verbs: ka = "go", tuʔ = "sleep" (no inflection seen) 4. CONSTRAINT IDENTIFICATION: - No agreement marking on verbs - Word order: PRONOUN + VERB 5. VERIFICATION: - Sentence 1: mo ka = "I go" - matches data - Sentence 2: ti ka = "you go" - matches data - Pattern is consistent 6. APPLICATION: - ti tuʔ = "you + sleep" = "you sleep" - mo ka = "I + go" = "I go" RULES: - RULE 1: mo -> "I" / _ (1st person singular pronoun) - RULE 2: ti -> "you" / _ (2nd person singular pronoun) - RULE 3: ka -> "go" / _ (uninflected verb) - RULE 4: tuʔ -> "sleep" / _ (uninflected verb) - RULE 5: [PRONOUN] + [VERB] -> [PRONOUN] + [VERB] EXPLANATION: The language uses independent pronouns and invariant verbs. Word order is PRONOUN-VERB. FINAL: ["you sleep", "I go"] === EXAMPLE 2: Morphological Rule Detection === DATA: 1. book | one book 2. books | two books 3. box | one box 4. boxes | two boxes QUERY: Complete the translation: 1. cat | one cat 2. cats | ___ 【GRAMMAR SCRATCHPAD】 1. DATA INVENTORY: - "book" appears as "book" (one) and "books" (two) - "box" appears as "box" (one) and "boxes" (two) - "one" marks singular, "two" marks plural 2. MORPHEME SEGMENTATION: - books = book + -s - boxes = box + -es - Singular forms lack suffix 3. PARADIGM MAPPING: - book ~ books: add -s - box ~ boxes: add -es - Pattern: plural = stem + suffix 4. CONSTRAINT IDENTIFICATION: - ALTERNATION: -s vs -es depends on stem-final sound - "book" ends in /k/ (non-sibilant) -> -s - "box" ends in /ks/ (sibilant) -> -es - RULE: Use -es after sibilants (/s/, /z/, /ʃ/, /tʃ/, /ks/, etc.) 5. VERIFICATION: - "book" + "s" = "books" - matches data - "box" + "es" = "boxes" - matches data - Rule holds 6. APPLICATION: - "cat" ends in /t/ (non-sibilant) -> plural = "cats" - "two" + "cats" = "two cats" RULES: - RULE 1: NOUN_sg -> NOUN_pl / [two] (context: plural number) - RULE 2: X -> X-s / _# (default plural: add -s) - RULE 3: X[sibilant] -> X-es / _# (sibilant plural: add -es) - RULE 4: Sibilant set: {s, z, ʃ, ʒ, tʃ, dʒ, ks, gz} - RULE 5: [two] + [NOUN_pl] -> "two" + [NOUN_pl] EXPLANATION: Plural formation uses -s by default. After sibilant sounds (s, z, sh, ch, x, etc.), use -es instead. "cat" ends in /t/, so regular -s plural applies. FINAL: ["one cat", "two cats"] """ GRAMMAR_SCRATCHPAD_TEMPLATE = """Before answering, you MUST work through this deduction process: 【GRAMMAR SCRATCHPAD】 1. DATA INVENTORY: List forms and glosses explicitly shown 2. MORPHEME SEGMENTATION: Break forms into morphemes, mark boundaries with - 3. PARADIGM MAPPING: Sketch inflectional categories (person, number, tense, case) 4. CONSTRAINT IDENTIFICATION: Phonological rules (sandhi, harmony, mutation), morphological rules (affixation, stem change) 5. VERIFICATION: Test hypothesis against exceptions 6. APPLICATION: Apply rules to query items step by step Then provide: RULES: List explicit rules in format "RULE N: [input] -> [output] / [condition]" EXPLANATION: Summary of grammatical analysis FINAL: JSON array with answers """ BASE_SYSTEM_PROMPT = ( "You are an expert linguist competing in the International Linguistics Olympiad. " "Analyze linguistic data, deduce grammar rules, and solve problems.\n\n" + FEW_SHOT_EXAMPLES + "\n\n" + GRAMMAR_SCRATCHPAD_TEMPLATE + "\n\n" "LINGUISTIC PHENOMENA TO CONSIDER:\n" "- Phonology: assimilation, dissimilation, deletion, epenthesis, metathesis, lenition, fortition\n" "- Morphology: prefixation, suffixation, infixation, circumfixation, reduplication, ablaut, suppletion\n" "- Phonotactics: consonant clusters, vowel constraints, syllable structure\n" "- Sandhi: external (word-boundary), internal (morpheme-boundary)\n" "- Syntax: word order (SOV, SVO, VSO), agreement patterns, case marking\n\n" "Respond with these sections:\n" "1. 【GRAMMAR SCRATCHPAD】 (complete all 6 steps with specific analysis)\n" "2. RULES: Explicit rewrite rules like:\n" " - RULE 1: [morpheme A] + [morpheme B] -> [result] / [environment]\n" " - RULE 2: X -> Y / _(condition) (phonological rule)\n" " - RULE 3: [CATEGORY] -> [translation/meaning]\n" "3. EXPLANATION: Clear summary of reasoning and rules\n" "4. FINAL: JSON array [\"answer1\", \"answer2\", ...] with exactly one string per query item" ) def build_router_aware_prompt(context: str, query: str, task_type: str, eval_type: str, k: int) -> str: """Build a prompt enriched with task routing information.""" profile = get_task_profile(task_type) header = f"""【TASK PROFILE】 Task Type: {task_type} Evaluation: {eval_type} (using {profile['eval_metric']}) Number of Items: {k} Focus: {profile['grammar_focus']} """ hint = profile["format_hint"] return f"{header}\n---\nDATA:\n{context.strip()}\n\nQUERY ({k} items):\n{query.strip()}\n\nFormat: {hint}" # ============================================================================= # 3. TEST-TIME SCALING (Self-Consistency Engine) # ============================================================================= def generate_n_paths(prompt: str, n: int = None, temperature: float = None, top_p: float = None) -> list: """Generate N diverse reasoning paths using sampling. FIX: now reads the (possibly tuned) global NUM_PATHS/TEMPERATURE/TOP_P/ MAX_NEW_TOKENS at call time via defaults, and top_p is actually threaded through to model.generate() instead of being hardcoded to 0.9. """ if n is None: n = NUM_PATHS if temperature is None: temperature = TEMPERATURE if top_p is None: top_p = TOP_P messages = [ {"role": "system", "content": BASE_SYSTEM_PROMPT}, {"role": "user", "content": prompt}, ] model_inputs = tok.apply_chat_template( messages, add_generation_prompt=True, return_tensors="pt" ) ids = model_inputs['input_ids'].to(model.device) paths = [] for path_idx in range(n): try: with torch.no_grad(): out = model.generate( ids, max_new_tokens=MAX_NEW_TOKENS, do_sample=True, temperature=temperature, top_p=top_p, pad_token_id=tok.eos_token_id, ) text = tok.decode(out[0][ids.shape[-1]:], skip_special_tokens=True).strip() paths.append(text) print(f" Path {path_idx+1}/{n} generated", flush=True) except Exception as e: print(f" Path {path_idx+1} failed: {e}", flush=True) paths.append("") return paths def chrF_score(candidate: str, reference: str) -> float: """Compute character n-gram F-score (simplified chrF).""" if not candidate or not reference: return 0.0 def get_ngrams(s, n): s = s.lower() return [s[i:i+n] for i in range(len(s) - n + 1)] total_f = 0.0 for n in range(2, 7): cand_ngrams = Counter(get_ngrams(candidate, n)) ref_ngrams = Counter(get_ngrams(reference, n)) overlapping = sum((cand_ngrams & ref_ngrams).values()) precision = overlapping / max(sum(cand_ngrams.values()), 1) recall = overlapping / max(sum(ref_ngrams.values()), 1) if precision + recall > 0: f = 2 * precision * recall / (precision + recall) total_f += f return total_f / 6.0 def consensus_vote(path_answers: list, eval_metric: str) -> tuple: """ Aggregate N paths using consensus voting. Returns (consensus_answer, confidence, explanation) """ if not path_answers or len(path_answers) == 0: return [], 0.0, "No paths generated" path_answers = [ans for ans in path_answers if ans] if not path_answers: return [], 0.0, "All paths failed" k = len(path_answers[0]) consensus = [] confidences = [] for item_idx in range(k): item_answers = [] for path in path_answers: if item_idx < len(path): item_answers.append(path[item_idx]) if not item_answers: consensus.append("") confidences.append(0.0) continue if eval_metric == "exact": answer_counts = Counter(item_answers) best_answer, count = answer_counts.most_common(1)[0] confidence = count / len(item_answers) consensus.append(best_answer) confidences.append(confidence) else: best_answer = item_answers[0] best_score = 0.0 for candidate in item_answers: score = sum(chrF_score(candidate, other) for other in item_answers) / len(item_answers) if score > best_score: best_score = score best_answer = candidate agreeing = sum(1 for ans in item_answers if chrF_score(best_answer, ans) > 0.8) confidence = agreeing / len(item_answers) consensus.append(best_answer) confidences.append(confidence) avg_confidence = sum(confidences) / len(confidences) if confidences else 0.0 explanation_parts = [ f"Generated {len(path_answers)} paths with T={TEMPERATURE}", f"Consensus ({eval_metric}) confidence: {avg_confidence:.0%}", ] return consensus, avg_confidence, " | ".join(explanation_parts) def extract_answers_from_text(text: str, expected_k: int) -> list: """Extract the JSON answer array from model output. FIX: the generic bracket-scanning fallback now scans matches in reverse (closest to the end of the text first), since FINAL: is always last in the expected output format, and earlier [...] occurrences (e.g. sibilant sets, category lists in RULES) were sometimes matched first and returned the wrong array when the FINAL: regex failed to match (e.g. truncated generation cut off before the closing bracket). """ # Try FINAL: block first m = re.search(r"FINAL:\s*(\[.*?\])", text, re.DOTALL | re.IGNORECASE) if m: try: parsed = json.loads(m.group(1)) if isinstance(parsed, list): return [str(x) for x in parsed] except Exception: pass # Look for any JSON array, preferring the one closest to the end of the # text (FINAL: is always the last section in the expected format). matches = list(re.finditer(r"\[.*?\]", text, re.DOTALL)) for m2 in reversed(matches): try: parsed = json.loads(m2.group(0)) if isinstance(parsed, list) and parsed: return [str(x) for x in parsed] except Exception: continue # Fallback: extract lines after FINAL: tail = text.split("FINAL:")[-1] lines = [re.sub(r"^\s*\d+[\.\)\]]\s*", "", l).strip(' \t\"\'') for l in tail.splitlines() if l.strip()] lines = [l for l in lines if l] return lines def extract_rules(text: str) -> list: """Extract explicit rules from RULES: section.""" rules = [] # Look for RULES: section rules_match = re.search(r"RULES:(.+?)(?=EXPLANATION:|FINAL:|$)", text, re.DOTALL | re.IGNORECASE) if rules_match: rules_text = rules_match.group(1) # Extract individual rule lines starting with - or RULE for line in rules_text.split('\n'): line = line.strip() if line and (line.startswith('-') or line.startswith('RULE') or (len(line) > 10 and '->' in line)): # Clean up the rule rule = re.sub(r'^[-•*]\s*', '', line).strip() if rule and len(rule) > 5: rules.append(rule) return rules def extract_explanation(text: str) -> str: """Extract the explanation portion from output.""" m = re.search(r"EXPLANATION:\s*(.+?)(?:RULES:|FINAL:|$)", text, re.DOTALL | re.IGNORECASE) if m: return m.group(1).strip()[:800] return "" # ============================================================================= # HYPERPARAMETER TUNING ENGINE # Optimizes: NUM_PATHS, TEMPERATURE, MAX_NEW_TOKENS, TOP_P # Uses small validation set to find best configuration # ============================================================================= HYPERPARAM_SPACE = { "num_paths": [2, 3, 4, 5], # N: number of reasoning paths "temperature": [0.1, 0.2, 0.3, 0.4, 0.5], # T: sampling diversity "top_p": [0.85, 0.9, 0.95, 0.99], # Nucleus sampling "max_tokens_factor": [1.0, 1.2, 1.5], # Multiplier for base tokens } # Validation set for hyperparameter tuning (subset of typical IOL problems) VALIDATION_EXAMPLES = [ { "id": "val-001", "task_type": "translation", "context": "1. áaka | I see\n2. tíika | you see\n3. áatʃi | I walk\n4. tíitʃi | you walk", "query": "Translate:\n1. áatʃi\n2. tíika", "expected": ["I walk", "you see"], "eval_metric": "chrF", }, { "id": "val-002", "task_type": "match_letters", "context": "1. atu A. water\n2. keno B. fire\n3. suna C. sun\n4. mizu D. east\n5. umi E. sea", "query": "Match:\n1. atu\n2. keno\n3. suna", "expected": ["A", "B", "C"], "eval_metric": "exact", }, ] def evaluate_hyperparams(config: dict, val_examples: list, max_evals: int = 2) -> dict: """ Evaluate a hyperparameter configuration on validation examples. Returns: {"score": float, "avg_time": float, "consistency": float} """ n_paths = config["num_paths"] temp = config["temperature"] top_p = config["top_p"] max_tokens = int(MAX_NEW_TOKENS * config["max_tokens_factor"]) total_score = 0.0 total_time = 0.0 consistency_scores = [] # Only evaluate on first max_evals examples for speed for example in val_examples[:max_evals]: start_t = time.time() # Generate N paths messages = [ {"role": "system", "content": BASE_SYSTEM_PROMPT}, {"role": "user", "content": example["context"] + "\n\n" + example["query"]}, ] model_inputs = tok.apply_chat_template( messages, add_generation_prompt=True, return_tensors="pt" ) ids = model_inputs['input_ids'].to(model.device) paths = [] for _ in range(n_paths): try: with torch.no_grad(): out = model.generate( ids, max_new_tokens=max_tokens, do_sample=True, temperature=temp, top_p=top_p, pad_token_id=tok.eos_token_id, ) text = tok.decode(out[0][ids.shape[-1]:], skip_special_tokens=True).strip() paths.append(text) except Exception: paths.append("") gen_time = time.time() - start_t total_time += gen_time # Extract answers path_answers = [] for path in paths: if path: answers = extract_answers_from_text(path, len(example["expected"])) answers = enforce_k_length(answers, len(example["expected"])) path_answers.append(answers) if not path_answers: continue # Calculate consensus and score if example["eval_metric"] == "exact": # For exact match: check if consensus matches expected answers_list = [tuple(a) for a in path_answers if a] if answers_list: most_common = Counter(answers_list).most_common(1)[0][0] consensus = list(most_common) score = sum(1 for i, exp in enumerate(example["expected"]) if i < len(consensus) and consensus[i] == exp) / len(example["expected"]) # Consistency = agreement among paths consistency = Counter(answers_list).most_common(1)[0][1] / len(answers_list) else: score = 0.0 consistency = 0.0 else: # For chrF: use fuzzy matching consensus, _, _ = consensus_vote(path_answers, "chrF") score = sum(chrF_score(consensus[i], example["expected"][i]) for i in range(min(len(consensus), len(example["expected"])))) / len(example["expected"]) # Consistency based on path agreement consistency = sum( sum(chrF_score(a1, a2) for a2 in path_answers) / len(path_answers) for a1 in path_answers ) / len(path_answers) if path_answers else 0.0 total_score += score consistency_scores.append(consistency) n_evaluated = min(max_evals, len(val_examples)) avg_score = total_score / n_evaluated if n_evaluated > 0 else 0.0 avg_time = total_time / n_evaluated if n_evaluated > 0 else 0.0 avg_consistency = sum(consistency_scores) / len(consistency_scores) if consistency_scores else 0.0 return { "score": avg_score, "avg_time": avg_time, "consistency": avg_consistency, } def grid_search_hyperparams(val_examples: list, max_configs: int = 8) -> dict: """ Grid search over hyperparameter space. Tests promising configurations and returns the best one. """ print("\n=== HYPERPARAMETER TUNING ===", flush=True) print(f"Testing configurations on {len(val_examples)} validation examples...", flush=True) # Priority order based on typical IOL performance test_configs = [ {"num_paths": 3, "temperature": 0.2, "top_p": 0.9, "max_tokens_factor": 1.0}, {"num_paths": 3, "temperature": 0.3, "top_p": 0.9, "max_tokens_factor": 1.0}, {"num_paths": 4, "temperature": 0.2, "top_p": 0.95, "max_tokens_factor": 1.2}, {"num_paths": 2, "temperature": 0.1, "top_p": 0.85, "max_tokens_factor": 1.0}, {"num_paths": 5, "temperature": 0.3, "top_p": 0.95, "max_tokens_factor": 1.2}, {"num_paths": 3, "temperature": 0.4, "top_p": 0.9, "max_tokens_factor": 1.0}, {"num_paths": 4, "temperature": 0.2, "top_p": 0.9, "max_tokens_factor": 1.5}, {"num_paths": 3, "temperature": 0.2, "top_p": 0.99, "max_tokens_factor": 1.2}, ][:max_configs] results = [] for i, config in enumerate(test_configs): # FIX: also bail out of tuning itself if time is getting short, so # HP search can't eat into the main processing budget unbounded. if time_left() < TIME_LIMIT * 0.6: print(f" Stopping HP search early: time_left={time_left():.0f}s", flush=True) break print(f"\nConfig {i+1}/{len(test_configs)}: N={config['num_paths']}, T={config['temperature']}, top_p={config['top_p']}", flush=True) metrics = evaluate_hyperparams(config, val_examples) results.append((config, metrics)) print(f" Score: {metrics['score']:.3f}, Consistency: {metrics['consistency']:.3f}, Time: {metrics['avg_time']:.2f}s", flush=True) if not results: # Nothing evaluated (ran out of time immediately) - fall back to defaults return test_configs[0] # Select best configuration based on combined score # Weight: accuracy 60%, consistency 30%, speed 10% def combined_score(config, metrics): # Normalize time (lower is better, assume max 30s per example) time_score = max(0, 1 - metrics['avg_time'] / 30.0) return (0.6 * metrics['score'] + 0.3 * metrics['consistency'] + 0.1 * time_score) best_config, best_metrics = max(results, key=lambda x: combined_score(x[0], x[1])) print(f"\n=== BEST CONFIGURATION ===", flush=True) print(f"NUM_PATHS: {best_config['num_paths']}", flush=True) print(f"TEMPERATURE: {best_config['temperature']}", flush=True) print(f"TOP_P: {best_config['top_p']}", flush=True) print(f"MAX_TOKENS_FACTOR: {best_config['max_tokens_factor']}", flush=True) print(f"Expected Score: {best_metrics['score']:.3f}", flush=True) return best_config # Global hyperparameters (will be set by tuning or use defaults) TUNED_NUM_PATHS = None TUNED_TEMPERATURE = None TUNED_TOP_P = None TUNED_MAX_TOKENS_FACTOR = None def apply_tuned_hyperparams(config: dict = None): """Apply tuned hyperparameters to global variables. FIX: previously TOP_P and the scaled MAX_NEW_TOKENS were computed and printed but never actually written back to the globals that generate_n_paths() reads, so all tuning was a no-op. Both are now applied. Also guards against re-scaling MAX_NEW_TOKENS more than once if this function is ever called twice. """ global NUM_PATHS, TEMPERATURE, TOP_P, MAX_NEW_TOKENS global TUNED_NUM_PATHS, TUNED_TEMPERATURE, TUNED_TOP_P, TUNED_MAX_TOKENS_FACTOR global _BASE_MAX_NEW_TOKENS if '_BASE_MAX_NEW_TOKENS' not in globals(): _BASE_MAX_NEW_TOKENS = MAX_NEW_TOKENS if config is None: # Use default/baseline configuration config = { "num_paths": NUM_PATHS, "temperature": TEMPERATURE, "top_p": TOP_P, "max_tokens_factor": 1.0, } TUNED_NUM_PATHS = config["num_paths"] TUNED_TEMPERATURE = config["temperature"] TUNED_TOP_P = config.get("top_p", TOP_P) TUNED_MAX_TOKENS_FACTOR = config.get("max_tokens_factor", 1.0) # Update the globals actually used by generate_n_paths() and friends NUM_PATHS = TUNED_NUM_PATHS TEMPERATURE = TUNED_TEMPERATURE TOP_P = TUNED_TOP_P MAX_NEW_TOKENS = int(_BASE_MAX_NEW_TOKENS * TUNED_MAX_TOKENS_FACTOR) print(f"\nApplied tuned hyperparameters:", flush=True) print(f" NUM_PATHS: {NUM_PATHS}", flush=True) print(f" TEMPERATURE: {TEMPERATURE}", flush=True) print(f" TOP_P: {TOP_P}", flush=True) print(f" MAX_NEW_TOKENS: {MAX_NEW_TOKENS}", flush=True) # ============================================================================= # REINFORCEMENT LEARNING COMPONENT # Online learning from test samples using self-consistency as reward # ============================================================================= class ReinforcementLearner: """ Online RL that samples from test data and reinforces successful patterns. Uses self-consistency and confidence as reward signals. """ def __init__(self): self.learned_patterns = {} # pattern -> success_score self.morpheme_inventory = {} # morpheme -> {translations, contexts} self.rule_confidence = {} # rule -> confidence_score self.paradigm_library = {} # pattern_type -> paradigms self.iteration_rewards = [] # Track reward over iterations def sample_test_data(self, df: pd.DataFrame, n_samples: int = 3) -> pd.DataFrame: """ Sample representative items from test set for RL training. Strategy: diverse task types, one (or more) sample per task type, trimmed to n_samples total. FIX: previous implementation concatenated per-task-type samples then applied .head(n_samples), which meant only the first task type(s) in groupby-iteration order ever survived the final truncation - i.e. it did NOT actually guarantee cross-task-type diversity as intended. This version samples a balanced number per group first, then trims with a random sample instead of a positional head(). """ if len(df) <= n_samples: return df.copy() groups = list(df.groupby('task_type')) if not groups: return df.sample(n_samples, random_state=0) per_group = max(1, n_samples // len(groups)) picked = [] for _, type_df in groups: type_df = type_df.copy() type_df['context_len'] = type_df['context'].str.len() median_len = type_df['context_len'].median() type_df['complexity_score'] = (type_df['context_len'] - median_len).abs() type_df = type_df.sort_values('complexity_score') picked.append(type_df.head(per_group)) sampled = pd.concat(picked) if len(sampled) > n_samples: sampled = sampled.sample(n_samples, random_state=0) return sampled.drop(columns=['context_len', 'complexity_score'], errors='ignore') def compute_reward(self, paths: list, consensus: list, confidence: float, expected_k: int, generation_time: float) -> dict: """ Compute multi-faceted reward for RL training. Reward components: - consistency_reward: Agreement between paths (self-consistency) - confidence_reward: Model confidence in its answer - format_reward: Proper JSON formatting and K-length - efficiency_reward: Reasonable generation time """ # Extract answers from all paths path_answers = [] for path in paths: if path: answers = extract_answers_from_text(path, expected_k) if answers: path_answers.append(tuple(answers)) # Consistency reward: pairwise agreement if len(path_answers) >= 2: agreements = 0 total_pairs = 0 for i in range(len(path_answers)): for j in range(i+1, len(path_answers)): # Compare answer similarity if len(path_answers[i]) == len(path_answers[j]): matches = sum(1 for a, b in zip(path_answers[i], path_answers[j]) if chrF_score(a, b) > 0.8) agreements += matches / len(path_answers[i]) total_pairs += 1 consistency_reward = agreements / total_pairs if total_pairs > 0 else 0.0 else: consistency_reward = 0.0 # Confidence reward (already computed) confidence_reward = confidence # Format reward: Did we get valid K-length outputs? format_reward = 1.0 if len(consensus) == expected_k and all(consensus) else 0.5 # Efficiency reward: Prefer faster generations (normalize to 30s target) efficiency_reward = max(0, 1 - generation_time / 30.0) # Combined reward (weighted) total_reward = ( 0.35 * consistency_reward + 0.35 * confidence_reward + 0.20 * format_reward + 0.10 * efficiency_reward ) return { "total": total_reward, "consistency": consistency_reward, "confidence": confidence_reward, "format": format_reward, "efficiency": efficiency_reward, } def extract_patterns(self, paths: list, consensus: list, task_type: str) -> list: """ Extract successful/reusable patterns from generated paths. """ patterns = [] for path in paths: if not path: continue # Extract morphemes from scratchpad morpheme_section = re.search( r'2\.\s*MORPHEME SEGMENTATION:(.+?)(?=3\.|PARADIGM|RULES|EXPLANATION|$)', path, re.DOTALL | re.IGNORECASE ) if morpheme_section: # Look for morpheme mappings like "word = morpheme1 + morpheme2" mappings = re.findall( r'([\w\-ʔʼʰʱʲʷː̥̩̪̯̃͡ɑ-ʯḀ-ẞ]+)\s*=\s*([\w\-ʔʼʰʱʲʷː̥̩̪̯̃͡ɑ-ʯḀ-ẞ]+)\s*\+\s*([\w\-ʔʼʰʱʲʷː̥̩̪̯̃͡ɑ-ʯḀ-ẞ]+)', morpheme_section.group(1) ) for mapping in mappings: pattern_key = f"SEG:{mapping[0]}->{mapping[1]}+{mapping[2]}" patterns.append(("morpheme_segmentation", pattern_key, mapping)) # Extract rules from RULES section rules = extract_rules(path) for rule in rules: rule_type = self._classify_rule(rule, task_type) patterns.append(("rule", rule_type, rule)) # Extract paradigm patterns paradigm_section = re.search( r'3\.\s*PARADIGM MAPPING:(.+?)(?=4\.|CONSTRAINT|RULES|EXPLANATION|$)', path, re.DOTALL | re.IGNORECASE ) if paradigm_section: # Look for category mappings categories = re.findall( r'(1st|2nd|3rd|singular|plural|present|past|nominative|accusative)\s*[=:]\s*([\w\-]+)', paradigm_section.group(1), re.IGNORECASE ) if categories: patterns.append(("paradigm", task_type, categories)) return patterns def _classify_rule(self, rule: str, task_type: str) -> str: """Classify rule type for organizing the library.""" rule_lower = rule.lower() if any(x in rule_lower for x in ['phonol', 'consonant', 'vowel', 'sound', 'harmony', 'sandhi']): return f"{task_type}:phonological" elif any(x in rule_lower for x in ['prefix', 'suffix', 'infix', 'affix', 'morpheme']): return f"{task_type}:morphological" elif any(x in rule_lower for x in ['word order', 'synta', 'sov', 'svo', 'agreement']): return f"{task_type}:syntactic" else: return f"{task_type}:general" def update_learned_patterns(self, patterns: list, reward: float): """ Update pattern library with exponential moving average. """ alpha = 0.3 # Learning rate for pattern_type, pattern_key, pattern_value in patterns: key = f"{pattern_type}:{pattern_key}" if key not in self.learned_patterns: self.learned_patterns[key] = { "type": pattern_type, "key": pattern_key, "value": pattern_value, "success_score": reward, "count": 1, } else: # Exponential moving average update old_score = self.learned_patterns[key]["success_score"] self.learned_patterns[key]["success_score"] = ( (1 - alpha) * old_score + alpha * reward ) self.learned_patterns[key]["count"] += 1 def build_rl_enhanced_prompt(self, base_prompt: str, task_type: str) -> str: """ Enhance prompt with learned patterns relevant to this task type. """ enhancements = [] # Add high-confidence patterns for this task type relevant_patterns = [ p for k, p in self.learned_patterns.items() if task_type in str(p.get("key", "")) and p["success_score"] > 0.5 ] # Sort by success score relevant_patterns.sort(key=lambda x: x["success_score"], reverse=True) if relevant_patterns: enhancements.append("\n【LEARNED PATTERNS FROM TEST DATA】") enhancements.append("Based on analysis of similar problems:") for i, pattern in enumerate(relevant_patterns[:3]): # Top 3 if pattern["type"] == "rule": enhancements.append(f" {i+1}. {pattern['key']}: {pattern['value']}") elif pattern["type"] == "morpheme_segmentation": val = pattern["value"] enhancements.append(f" {i+1}. Morpheme pattern: {val[0]} = {val[1]} + {val[2]}") enhancements.append("Consider these patterns in your analysis.\n") return base_prompt + "\n" + "\n".join(enhancements) if enhancements else base_prompt def rl_training_loop(self, df: pd.DataFrame, n_iterations: int = 2, samples_per_iter: int = 3, time_cap_seconds: float = None) -> dict: """ Main RL training loop: sample, evaluate, reinforce, repeat. Returns optimized hyperparameters and learned patterns. FIX: previously the only time check inside the loop happened once per sample row (before starting a config's generation batch), so a single slow batch of generations could still blow well past the intended budget and starve the main per-row loop of time, resulting in a submission full of empty timeout-fallback predictions. Now there is a hard wall-clock deadline (time_cap_seconds, relative to START) that is checked before every single model.generate() call, not just once per sample/config. """ print(f"\n=== REINFORCEMENT LEARNING ({n_iterations} iterations) ===", flush=True) if time_cap_seconds is None: time_cap_seconds = TIME_LIMIT * 0.15 # hard default cap: 15% of total budget rl_deadline = START + time_cap_seconds best_config = None best_avg_reward = 0.0 # Configurations to try (will be refined based on rewards) configs_to_try = [ {"num_paths": 3, "temperature": 0.2, "top_p": 0.9}, {"num_paths": 4, "temperature": 0.3, "top_p": 0.95}, {"num_paths": 5, "temperature": 0.25, "top_p": 0.92}, ] for iteration in range(n_iterations): if time.time() > rl_deadline: print(f" RL time cap reached before iteration {iteration+1}, stopping", flush=True) break print(f"\n--- RL Iteration {iteration + 1}/{n_iterations} ---", flush=True) # Sample test data sample_df = self.sample_test_data(df, samples_per_iter) print(f"Sampled {len(sample_df)} items for training", flush=True) iteration_rewards = [] for _, row in sample_df.iterrows(): if time.time() > rl_deadline: print(" RL time cap reached mid-iteration, stopping", flush=True) break task_type = row.get("task_type", "translation") eval_type = row.get("eval_type", "chr_f1") eval_metric = "exact" if eval_type.startswith("exact") else "chrF" k = count_items(row["query"]) # Build base prompt base_prompt = build_router_aware_prompt( row["context"], row["query"], task_type, eval_type, k ) # Enhance with learned patterns enhanced_prompt = self.build_rl_enhanced_prompt(base_prompt, task_type) # Try different configurations for config in configs_to_try[:2]: # Try first 2 configs if time.time() > rl_deadline: print(" RL time cap reached before config trial, stopping", flush=True) break start_t = time.time() # Generate paths messages = [ {"role": "system", "content": BASE_SYSTEM_PROMPT}, {"role": "user", "content": enhanced_prompt}, ] model_inputs = tok.apply_chat_template( messages, add_generation_prompt=True, return_tensors="pt" ) ids = model_inputs['input_ids'].to(model.device) # Also respect the absolute competition time limit if time_left() < TIME_LIMIT * 0.4: print(f" RL sample aborted: critical time ({time_left():.0f}s)", flush=True) break paths = [] for _ in range(config["num_paths"]): # FIX: check the deadline before every single generation, # not just once per sample/config. if time.time() > rl_deadline or time_left() < TIME_LIMIT * 0.35: print(" RL time cap reached mid-generation, stopping this batch", flush=True) break try: with torch.no_grad(): out = model.generate( ids, max_new_tokens=MAX_NEW_TOKENS, do_sample=True, temperature=config["temperature"], top_p=config["top_p"], pad_token_id=tok.eos_token_id, ) text = tok.decode(out[0][ids.shape[-1]:], skip_special_tokens=True).strip() paths.append(text) except Exception: paths.append("") gen_time = time.time() - start_t if not paths: continue # Extract answers and compute consensus path_answers = [] for path in paths: if path: ans = extract_answers_from_text(path, k) path_answers.append(enforce_k_length(ans, k)) if path_answers: consensus, confidence, _ = consensus_vote(path_answers, eval_metric) consensus = enforce_k_length(consensus, k) # Compute reward reward = self.compute_reward( paths, consensus, confidence, k, gen_time ) iteration_rewards.append(reward["total"]) # Extract and reinforce patterns if reward is good if reward["total"] > 0.5: patterns = self.extract_patterns(paths, consensus, task_type) self.update_learned_patterns(patterns, reward["total"]) # Track best config if reward["total"] > best_avg_reward: best_avg_reward = reward["total"] best_config = config avg_reward = sum(iteration_rewards) / len(iteration_rewards) if iteration_rewards else 0.0 self.iteration_rewards.append(avg_reward) print(f" Average reward: {avg_reward:.3f}", flush=True) print(f" Learned patterns: {len(self.learned_patterns)}", flush=True) # Determine final config if best_config is None: best_config = configs_to_try[0] # Default print(f"\n=== RL TRAINING COMPLETE ===", flush=True) print(f"Best config: N={best_config['num_paths']}, T={best_config['temperature']}", flush=True) print(f"Best reward: {best_avg_reward:.3f}", flush=True) print(f"Total learned patterns: {len(self.learned_patterns)}", flush=True) print(f"Time spent in RL: {time.time() - START:.1f}s (cap was {time_cap_seconds:.0f}s)", flush=True) return { "best_config": best_config, "learned_patterns": self.learned_patterns, "reward_history": self.iteration_rewards, } def get_learned_pattern_summary(self) -> str: """Return summary of learned patterns for explanation.""" if not self.learned_patterns: return "No patterns learned yet" top_patterns = sorted( self.learned_patterns.values(), key=lambda x: x["success_score"], reverse=True )[:5] summary = [] for p in top_patterns: summary.append(f" - {p['key']} (score: {p['success_score']:.2f}, n={p['count']})") return "\n".join(summary) # Global RL learner instance rl_learner = ReinforcementLearner() # ============================================================================= # 4. DETERMINISTIC ALIGNMENT & FALLBACK GUARDRAIL # ============================================================================= def enforce_k_length(answers: list, k: int) -> list: """Ensure output is exactly K items.""" answers = list(answers)[:k] while len(answers) < k: answers.append("") return answers def chrF_fallback_recovery(paths: list, expected_k: int, target_item_idx: int) -> str: """When consensus fails, try chrF-based selection.""" candidates = [] for path in paths: answers = extract_answers_from_text(path, expected_k) if target_item_idx < len(answers) and answers[target_item_idx]: candidates.append(answers[target_item_idx]) if not candidates: return "" counts = Counter(candidates) return counts.most_common(1)[0][0] def process_row_with_architecture(row: pd.Series, use_rl_prompt: bool = True) -> dict: """Main processing function with optional RL enhancement.""" task_type = row.get("task_type", "translation") eval_type = row.get("eval_type", "chr_f1") context = row["context"] query = row["query"] k = count_items(query) profile = get_task_profile(task_type) eval_metric = "exact" if eval_type.startswith("exact") else "chrF" print(f"\n[{row['id']}] Task={task_type}, Eval={eval_metric}, K={k}", flush=True) base_prompt = build_router_aware_prompt(context, query, task_type, eval_type, k) # Enhance with learned patterns if RL is active if use_rl_prompt: prompt = rl_learner.build_rl_enhanced_prompt(base_prompt, task_type) else: prompt = base_prompt print(f" Generating N={NUM_PATHS} paths with Few-Shot Rules...", flush=True) paths = generate_n_paths(prompt, n=NUM_PATHS, temperature=TEMPERATURE, top_p=TOP_P) path_answers = [] all_rules = [] for i, path in enumerate(paths): if path: answers = extract_answers_from_text(path, k) answers = enforce_k_length(answers, k) path_answers.append(answers) rules = extract_rules(path) all_rules.extend(rules) print(f" Path {i+1}: {answers}", flush=True) if path_answers: consensus, confidence, voting_explanation = consensus_vote(path_answers, eval_metric) consensus = enforce_k_length(consensus, k) print(f" Consensus: {consensus} (confidence={confidence:.2f})", flush=True) else: consensus = [""] * k confidence = 0.0 voting_explanation = "No valid paths" final_answers = enforce_k_length(consensus, k) if confidence < 0.5 and path_answers: for i in range(k): if not final_answers[i] or not final_answers[i].strip(): recovered = chrF_fallback_recovery(paths, k, i) if recovered: final_answers[i] = recovered print(f" chrF recovery for item {i+1}: {recovered}", flush=True) final_answers = enforce_k_length(final_answers, k) best_explanation = "" for path in paths: if path: best_explanation = extract_explanation(path) if best_explanation: break unique_rules = list(dict.fromkeys([r for r in all_rules if r])) rules_summary = " | ".join(unique_rules[:5]) if unique_rules else "No rules" full_explanation = f"{voting_explanation} | {best_explanation[:350]} | Rules: {rules_summary[:250]}" return { "id": row["id"], "pred": json.dumps(final_answers, ensure_ascii=False), "explanation": full_explanation[:1200], } # ============================================================================= # SETUP: Reinforcement Learning + Hyperparameter Tuning # # FIX: both phases are now hard-capped in wall-clock time and gated by much # more conservative time-remaining thresholds, so neither can meaningfully # eat into the budget needed for the main per-row processing loop. RL is # capped at 15% of TIME_LIMIT; HP tuning only runs if RL did not run/found # nothing, and internally bails once time_left() drops below 60% of budget. # ============================================================================= rl_results = None # Phase 1: Reinforcement Learning from test samples (if time permits and not disabled) if not DISABLE_RL and time_left() > TIME_LIMIT * 0.85: # only with very large headroom try: print("\n" + "="*60, flush=True) print("PHASE 1: REINFORCEMENT LEARNING FROM TEST DATA", flush=True) print("="*60, flush=True) rl_results = rl_learner.rl_training_loop( df, n_iterations=2, samples_per_iter=3, time_cap_seconds=TIME_LIMIT * 0.15, ) # Apply RL-discovered config if rl_results and rl_results.get("best_config"): cfg = rl_results["best_config"] NUM_PATHS = cfg["num_paths"] TEMPERATURE = cfg["temperature"] TOP_P = cfg.get("top_p", TOP_P) print(f"\n>>> Applied RL config: N={NUM_PATHS}, T={TEMPERATURE}, top_p={TOP_P}", flush=True) print(f"\n>>> Learned patterns:") print(rl_learner.get_learned_pattern_summary(), flush=True) except Exception as e: print(f"RL training error: {e}", flush=True) # Phase 2: Traditional hyperparameter tuning (if RL didn't run or failed) if not rl_results and time_left() > TIME_LIMIT * 0.7: try: print("\n" + "="*60, flush=True) print("PHASE 2: HYPERPARAMETER TUNING", flush=True) print("="*60, flush=True) best_config = grid_search_hyperparams(VALIDATION_EXAMPLES, max_configs=4) apply_tuned_hyperparams(best_config) except Exception as e: print(f"HP tuning error: {e}", flush=True) apply_tuned_hyperparams(None) else: if not rl_results: print("Skipping tuning - time constrained", flush=True) apply_tuned_hyperparams(None) print(f"\n>>> Time remaining before main loop: {time_left():.0f}s / {TIME_LIMIT}s", flush=True) # ============================================================================= # MAIN PROCESSING LOOP (with RL-enhanced prompts) # ============================================================================= rows_out = [] n_rows = len(df) # Main processing with time awareness and graceful degradation for i, r in df.iterrows(): expected_k = count_items(r["query"]) remaining_rows = n_rows - i per_row_budget = (time_left() - SAFETY_BUFFER) / max(1, remaining_rows) # Safety: absolute time cutoff if time_left() < SAFETY_BUFFER: print(f"TIMEOUT: Only {time_left():.0f}s left, using fast fallback", flush=True) rows_out.append({ "id": r["id"], "pred": json.dumps([""] * expected_k, ensure_ascii=False), "explanation": "Timeout: processing stopped", }) continue if time_left() < SAFETY_BUFFER or per_row_budget < 10: rows_out.append({ "id": r["id"], "pred": json.dumps([""] * expected_k, ensure_ascii=False), "explanation": "Timeout fallback", }) continue try: result = process_row_with_architecture(r) rows_out.append(result) except Exception as e: print(f"Error processing {r['id']}: {e}", flush=True) rows_out.append({ "id": r["id"], "pred": json.dumps([""] * expected_k, ensure_ascii=False), "explanation": f"Error: {str(e)[:150]}", }) pd.DataFrame(rows_out).to_csv("submission.csv", index=False) print(f"Progress: {len(rows_out)}/{n_rows} rows, {time_left():.0f}s left", flush=True) pd.DataFrame(rows_out, columns=["id", "pred", "explanation"]).to_csv( "submission.csv", index=False) print("wrote submission.csv", flush=True) print(f"Total time: {time.time() - START:.1f}s", flush=True)