1342 lines
54 KiB
Python
1342 lines
54 KiB
Python
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 = 4096 # increased for scratchpad + rules
|
||
NUM_PATHS = 3 # N=3 for test-time scaling
|
||
TEMPERATURE = 0.1 # 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) |