Upload folder using huggingface_hub

script.py

@@ -1,4 +1,5 @@
 
+
 import os
 os.environ["HF_HUB_OFFLINE"] = "1"
 os.environ["TRANSFORMERS_OFFLINE"] = "1"
@@ -20,10 +21,6 @@ try:
 except ImportError:
     _run_pkg_cmd(["install", "-q", "bitsandbytes"])
 
-# torchvision is not needed for a text-only LLM, and if the sandbox's
-# preinstalled torchvision doesn't match its torch build, transformers can
-# crash on import trying to load an unused image-utils path. Drop it
-# defensively before importing transformers.
 try:
     import torchvision  # noqa: F401
     _run_pkg_cmd(["uninstall", "-y", "-q", "torchvision"])
@@ -33,14 +30,18 @@ except ImportError:
 import re
 import json
 import time
+import math
 import pandas as pd
 import torch
+from collections import Counter
 from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 
 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 = 900
+MAX_NEW_TOKENS = 1200         # increased for grammar scratchpad
+NUM_PATHS = 3                 # N=3 for test-time scaling
+TEMPERATURE = 0.3             # T=0.3 for diverse but coherent paths
 
 START = time.time()
 
@@ -53,7 +54,7 @@ tok = AutoTokenizer.from_pretrained(MODEL_ID)
 bnb_config = BitsAndBytesConfig(
     load_in_4bit=True,
     bnb_4bit_quant_type="nf4",
-    bnb_4bit_compute_dtype=torch.float16,   # T4 = Turing, no native bf16
+    bnb_4bit_compute_dtype=torch.float16,
     bnb_4bit_use_double_quant=True,
 )
 
@@ -69,43 +70,226 @@ df = pd.read_csv("/tmp/data/test.csv", dtype=str).fillna("")
 
 ITEM_RE = re.compile(r"(?m)^\s*(\d+)\.\s")
 
-TASK_HINTS = {
-    "translation": "Answer with the translated text only, in the language the item asks for.",
-    "match_letters": "Answer with only the option letter (e.g. A, B, C) for each item.",
-    "fill_blanks": "Answer with only the missing form for each blank, using the same spelling/diacritic conventions as the data.",
-    "text_to_num": "Answer with the value written in digits (e.g. 285), no words.",
-    "num_to_text": "Answer with the number written out in words in the task language, matching the data's conventions.",
-}
+# =============================================================================
+# 1. DYNAMIC TASK ROUTER
+# Analyzes: task_type, eval_type, item count (K)
+# =============================================================================
 
-SYSTEM_PROMPT = (
-    "You are a world-class contestant at the International Linguistics Olympiad (IOL). "
-    "You will be given a self-contained problem: some linguistic data plus a query with "
-    "numbered items. Work out the pattern rigorously from the data given -- do not rely on "
-    "any outside knowledge of the language, only on what is in the context. Think step by "
-    "step internally, but keep your visible reasoning compact.\n\n"
-    "You MUST end your reply with exactly two blocks, in this order:\n\n"
-    "EXPLANATION:\n"
-    "<a short, human-readable explanation of your reasoning and the rule(s) you found; "
-    "a few bullet points or a couple of short paragraphs, NOT a raw token-by-token trace>\n\n"
-    "FINAL:\n"
-    "<a JSON list of strings, one answer per numbered item, in the same order as the query, "
-    "and nothing else on that line>"
-)
+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 n_items(query: str) -> int:
+def count_items(query: str) -> int:
+    """Extract K = number of items in the query."""
     nums = ITEM_RE.findall(query)
     if nums:
         return len(set(nums))
-    # fallback: count non-empty lines as a rough guess
+    # fallback: count non-empty lines
     return max(1, len([l for l in query.splitlines() if l.strip()]))
 
-def build_prompt(context: str, query: str, task_type: str) -> str:
-    hint = TASK_HINTS.get(task_type, "")
-    extra = f"\n\nFormat hint: {hint}" if hint else ""
-    return f"{context.strip()}\n\n{query.strip()}{extra}"
+def get_task_profile(task_type: str) -> dict:
+    """Route to task-specific configuration."""
+    return TASK_PROFILES.get(task_type, TASK_PROFILES["translation"])
+
+# =============================================================================
+# 2. THE GRAMMAR SCRATCHPAD (System Prompt Injection)
+# Forces syntax deduction before answering
+# =============================================================================
+
+GRAMMAR_SCRATCHPAD = """Before answering, you MUST work through this deduction process in your scratchpad:
+
+【GRAMMAR SCRATCHPAD】
+1. DATA INVENTORY: List the forms and glosses explicitly shown in the data
+2. MORPHEME SEGMENTATION: Break forms into plausible morpheme boundaries
+3. PARADIGM MAPPING: If there are inflections, sketch the paradigm (person, number, tense, case...)
+4. CONSTRAINT IDENTIFICATION: What rules govern the forms? (sandhi, vowel harmony, consonant mutation...)
+5. VERIFICATION: Test your hypothesis against any exceptions in the data
+6. APPLICATION: Apply the deduced rules to the query items
+
+After completing the scratchpad, provide your FINAL answer.
+"""
+
+BASE_SYSTEM_PROMPT = (
+    "You are an expert linguist competing in the International Linguistics Olympiad. "
+    "You will analyze linguistic data, deduce the underlying grammar rules, and solve problems.\n\n"
+    + GRAMMAR_SCRATCHPAD + "\n\n"
+    "Respond with:\n"
+    "1. EXPLANATION: A clear summary of your reasoning and the rules you found\n"
+    "2. FINAL: A JSON array with exactly one answer string per numbered item\n"
+    "   - Preserve diacritics and special characters exactly\n"
+    "   - Match the orthography seen in the examples\n"
+    "   - Do not include any markdown formatting or extra text"
+)
+
+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)
+# Generates N=3 paths (T=0.3) -> Consensus Voting
+# =============================================================================
+
+def generate_n_paths(prompt: str, n: int = NUM_PATHS, temperature: float = TEMPERATURE) -> list:
+    """Generate N diverse reasoning paths using sampling."""
+    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=0.9,
+                    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("")  # placeholder for failed generation
+    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)]
+
+    # Use n-grams of size 2, 3, 4, 5, 6 (chrF default)
+    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  # average over n-gram sizes
+
+def consensus_vote(path_answers: list, eval_metric: str) -> tuple:
+    """
+    Aggregate N paths using consensus voting.
+    - For exact match: majority vote (string equality)
+    - For chrF: fuzzy clustering and highest-confidence cluster
+    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]  # filter empty
+    if not path_answers:
+        return [], 0.0, "All paths failed"
+
+    k = len(path_answers[0])
+    consensus = []
+    confidences = []
+
+    for item_idx in range(k):
+        # Extract item answer from each path
+        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":
+            # Exact match majority voting
+            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:
+            # chrF-based clustering (fuzzy voting)
+            # Find the answer that maximizes average similarity to others
+            best_answer = item_answers[0]
+            best_score = 0.0
 
-def extract_final_list(text: str, expected_n: int):
-    m = re.search(r"FINAL:\s*(\[.*?\])", text, re.DOTALL)
+            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
+
+            # Confidence based on agreement
+            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
+
+    # Build explanation showing voting outcome
+    explanation_parts = [
+        f"Generated {len(path_answers)} reasoning paths with T={TEMPERATURE}.",
+        f"Consensus voting ({eval_metric} metric) selected answers with {avg_confidence:.0%} avg confidence.",
+    ]
+
+    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."""
+    # Try FINAL: block first
+    m = re.search(r"FINAL:\s*(\[.*?\])", text, re.DOTALL | re.IGNORECASE)
     if m:
         try:
             parsed = json.loads(m.group(1))
@@ -113,7 +297,8 @@ def extract_final_list(text: str, expected_n: int):
                 return [str(x) for x in parsed]
         except Exception:
             pass
-    # fallback 1: any JSON array anywhere in the text
+
+    # Look for any JSON array
     for m2 in re.finditer(r"\[.*?\]", text, re.DOTALL):
         try:
             parsed = json.loads(m2.group(0))
@@ -121,86 +306,166 @@ def extract_final_list(text: str, expected_n: int):
                 return [str(x) for x in parsed]
         except Exception:
             continue
-    # fallback 2: split remaining non-empty lines after FINAL: (or whole text)
+
+    # 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 = [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]
-    if lines:
-        return lines
-    return []
+    return lines
 
 def extract_explanation(text: str) -> str:
-    m = re.search(r"EXPLANATION:\s*(.*?)\s*FINAL:", text, re.DOTALL)
-    expl = m.group(1).strip() if m else ""
-    return expl[:1200]
-
-def fit_to_length(items, n):
-    items = list(items)
-    if len(items) < n:
-        items = items + [""] * (n - len(items))
-    elif len(items) > n:
-        items = items[:n]
-    return items
+    """Extract the explanation portion from output."""
+    m = re.search(r"EXPLANATION:\s*(.*?)\s*(?:FINAL:|$)", text, re.DOTALL | re.IGNORECASE)
+    if m:
+        return m.group(1).strip()[:1000]
+    return ""
+
+# =============================================================================
+# 4. DETERMINISTIC ALIGNMENT & FALLBACK GUARDRAIL
+# Enforces exact K-length JSON array & chrF recovery
+# =============================================================================
+
+def enforce_k_length(answers: list, k: int, context: str = "", query: str = "") -> list:
+    """
+    Deterministically ensure output is exactly K items.
+    - Pad with empty strings if too few
+    - Truncate if too many
+    """
+    answers = list(answers)[:k]  # Take first k
+    while len(answers) < k:
+        answers.append("")  # Pad to k
+    return answers
+
+def chrF_fallback_recovery(paths: list, expected_k: int, target_item_idx: int, 
+                           reference: str = "") -> str:
+    """
+    When consensus fails, try chrF-based selection from all paths.
+    """
+    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 ""
+
+    # Return the answer that appears most frequently (fallback voting)
+    counts = Counter(candidates)
+    return counts.most_common(1)[0][0]
+
+def process_row_with_architecture(row: pd.Series) -> dict:
+    """
+    Main processing function implementing the 4-stage architecture.
+    """
+    # STAGE 1: Dynamic Task Router
+    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)
+
+    # Build router-aware prompt with grammar scratchpad hint
+    prompt = build_router_aware_prompt(context, query, task_type, eval_type, k)
+
+    # STAGE 2 & 3: Grammar Scratchpad + Test-Time Scaling
+    print(f"  Generating N={NUM_PATHS} paths with Grammar Scratchpad...", flush=True)
+    paths = generate_n_paths(prompt, n=NUM_PATHS, temperature=TEMPERATURE)
+
+    # Extract answers from each path
+    path_answers = []
+    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)
+            print(f"    Path {i+1}: {answers}", flush=True)
+
+    # STAGE 3: Consensus Voting
+    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 generated"
+
+    # STAGE 4: Deterministic Alignment & Fallback
+    final_answers = enforce_k_length(consensus, k)
+
+    # chrF fallback: if confidence is low, try to recover
+    if confidence < 0.5 and path_answers:
+        for i in range(k):
+            if not final_answers[i] or 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 alignment check
+    final_answers = enforce_k_length(final_answers, k)
+
+    # Extract explanation from best path (highest confidence answer source)
+    best_explanation = ""
+    for path in paths:
+        if path:
+            best_explanation = extract_explanation(path)
+            if best_explanation:
+                break
+
+    # Combine with voting explanation
+    full_explanation = f"{voting_explanation} | Grammar Analysis: {best_explanation[:500]}"
+
+    return {
+        "id": row["id"],
+        "prediction": json.dumps(final_answers, ensure_ascii=False),
+        "explanation": full_explanation[:1200],
+    }
+
+# =============================================================================
+# MAIN PROCESSING LOOP
+# =============================================================================
 
 rows_out = []
 n_rows = len(df)
 
 for i, r in df.iterrows():
-    expected_n = n_items(r["query"])
+    expected_k = count_items(r["query"])
     remaining_rows = n_rows - i
     per_row_budget = (time_left() - SAFETY_BUFFER) / max(1, remaining_rows)
 
-    if time_left() < SAFETY_BUFFER or per_row_budget < 5:
-        # out of time: best-effort blank guesses so every row still has a row
+    if time_left() < SAFETY_BUFFER or per_row_budget < 10:
+        # Out of time: use fast fallback
         rows_out.append({
             "id": r["id"],
-            "prediction": json.dumps([""] * expected_n, ensure_ascii=False),
-            "explanation": "",
+            "prediction": json.dumps([""] * expected_k, ensure_ascii=False),
+            "explanation": "Timeout fallback - no computation time remaining",
         })
         continue
 
     try:
-        prompt = build_prompt(r["context"], r["query"], r.get("task_type", ""))
-        messages = [
-            {"role": "system", "content": 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)
-
-        gen_tokens = min(MAX_NEW_TOKENS, MAX_NEW_TOKENS)
-        with torch.no_grad():
-            out = model.generate(
-                ids,
-                max_new_tokens=gen_tokens,
-                do_sample=False,
-                temperature=None,
-                top_p=None,
-                pad_token_id=tok.eos_token_id,
-            )
-        text = tok.decode(out[0][ids.shape[-1]:], skip_special_tokens=True).strip()
-
-        preds = extract_final_list(text, expected_n)
-        preds = fit_to_length(preds, expected_n)
-        explanation = extract_explanation(text)
-
-        rows_out.append({
-            "id": r["id"],
-            "prediction": json.dumps(preds, ensure_ascii=False),
-            "explanation": explanation,
-        })
+        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"],
-            "prediction": json.dumps([""] * expected_n, ensure_ascii=False),
-            "explanation": f"(error during generation: {e})"[:300],
+            "prediction": json.dumps([""] * expected_k, ensure_ascii=False),
+            "explanation": f"Error: {str(e)[:200]}",
         })
 
-    # write after every row so a partial file is always valid on disk
+    # Write after every row for persistence
     pd.DataFrame(rows_out).to_csv("submission.csv", index=False)
-    print(f"{len(rows_out)}/{n_rows} done, {time_left():.0f}s left", flush=True)
+    print(f"Progress: {len(rows_out)}/{n_rows} rows, {time_left():.0f}s left", flush=True)
 
+# Final write
 pd.DataFrame(rows_out).to_csv("submission.csv", index=False)
-print("wrote submission.csv", flush=True)
+print("\n=== submission.csv written successfully ===", flush=True)
+print(f"Total time: {time.time() - START:.1f}s", flush=True)