script.py

# =============================================================================
# DIAGNOSTIC v2: Information compression vs one-pass baseline
#
# CORRECTED HYPOTHESIS:
# Does forcing the model to extract PROVABLE OBSERVATIONS from the examples,
# then answer using those observations + examples, improve EM over one-pass?
#
# NOT: "Can Qwen become a linguist?" (grammar induction, speculative)
# YES: "Does organizing evidence improve exact reasoning?" (compression, verifiable)
#
# CONTROLS FOR CONFOUNDS:
#
# 1. TWO-CALL CONFOUND (extra compute vs mechanism):
#    Three arms, not two:
#      A. Baseline (1 call, greedy)
#      B. Baseline with extra reasoning tokens (1 call, extra budget) -- CONTROL
#      C. Observations + Answer (2 calls, greedy) -- TREATMENT
#    If B >= C, the gain is compute, not the mechanism.
#
# 2. NON-REPRODUCIBILITY:
#    All three arms use do_sample=False. Fully deterministic.
#    We test the MECHANISM (does compression help?), not sampling variance.
#
# 3. HALLUCINATION:
#    Observation prompt forbids inference beyond what examples demonstrate.
#    Each observation must cite which examples support it.
#    Observations that don't cite are rejected.
#
# 4. TASK LEAKAGE:
#    Observation pass does NOT see the query.
#    It sees ONLY the examples and the general task type.
#    This measures "examples -> observations", not "examples + task hint -> answer prep".
#
# 5. RUNTIME:
#    Observation pass is capped at 256 tokens (shorter than answer generation).
#    Total per problem: ~ baseline_time * 1.5, not 2x.
#
# Run in Colab after loading model + Linguini df with columns:
#   context, query, answer, task_type
# =============================================================================

import re, unicodedata, time
import torch

# ---------------------------------------------------------------------------
# Frozen normalization from proven 0.121 baseline
# ---------------------------------------------------------------------------
def norm_generic(s):
    s = unicodedata.normalize("NFC", s).strip()
    s = re.sub(r"^\s*\(?\d+\)?\s*[.):\-]\s+", "", s)
    s = re.sub(r"^\s*[-*•·]\s+", "", s)
    s = re.sub(r"(?i)^\s*(answer|translation|output)\s*\d*\s*[:.\-]\s+", "", s)
    s = s.strip("* ")
    s = re.sub(r"\s{2,}", " ", s)
    return s.strip()

def norm_number(s):
    g = norm_generic(s)
    m = re.search(r"-?\d[\d,\. ]*\d|\d", g)
    if not m: return g
    digits = re.sub(r"[^\d-]", "", m.group(0))
    return digits if digits else g

def normalize(task_type, s):
    if task_type == "text_to_num":
        return norm_number(s)
    return norm_generic(s)

def n_expected(query):
    items = re.findall(r"(?m)^\s*(\d+)\s*[.\)]", query)
    if items: return len(items)
    rng = re.search(r"\(\s*(\d+)\s*[-–—]\s*(\d+)\s*\)", query)
    if rng:
        lo, hi = int(rng.group(1)), int(rng.group(2))
        if 0 < hi - lo < 100: return hi - lo + 1
    return None

# ---------------------------------------------------------------------------
# Deterministic greedy generation -- one function, three arms use identical
# calling convention except for max_new_tokens
# ---------------------------------------------------------------------------
def generate(messages, max_new_tokens=512):
    try:
        enc = tok.apply_chat_template(messages, add_generation_prompt=True,
            enable_thinking=False, return_tensors="pt", return_dict=True).to(model.device)
        ilen = enc["input_ids"].shape[-1]
        with torch.no_grad():
            out = model.generate(**enc, max_new_tokens=max_new_tokens, do_sample=False)
    except Exception:
        ids = tok.apply_chat_template(messages, add_generation_prompt=True,
            enable_thinking=False, return_tensors="pt").to(model.device)
        ilen = ids.shape[-1]
        with torch.no_grad():
            out = model.generate(ids, max_new_tokens=max_new_tokens, do_sample=False)
    return tok.decode(out[0][ilen:], skip_special_tokens=True).strip()

# ---------------------------------------------------------------------------
# ARM A: Baseline (identical to proven 0.121 script)
# ---------------------------------------------------------------------------
BASELINE_SYS = (
    "You solve International Linguistics Olympiad problems about a language you have never seen. "
    "Everything you need is in the examples. Answer every numbered item, in order. "
    "Put each answer on its own line, with no numbering and no extra text."
)

def arm_baseline(context, query):
    msgs = [
        {"role": "system", "content": BASELINE_SYS},
        {"role": "user",   "content": f"{context.strip()}\n\n{query.strip()}"}
    ]
    return generate(msgs, max_new_tokens=512)

# ---------------------------------------------------------------------------
# ARM B: Baseline + extra reasoning budget (CONTROL for compute)
# Same single call, same greedy, same prompt -- only difference is more tokens.
# If this arm matches Arm C, the gain is compute, not the mechanism.
# ---------------------------------------------------------------------------
def arm_baseline_extended(context, query):
    msgs = [
        {"role": "system", "content": BASELINE_SYS},
        {"role": "user",   "content": f"{context.strip()}\n\n{query.strip()}"}
    ]
    # Same prompt, more tokens -- controls for extra compute
    return generate(msgs, max_new_tokens=768)

# ---------------------------------------------------------------------------
# ARM C: Observations then Answer
# ---------------------------------------------------------------------------
# OBSERVATION PROMPT: strict grounding, no task leakage
# - Does NOT see the query
# - Every observation must cite supporting examples
# - Forbids inference beyond what examples demonstrate
OBSERVATION_SYS = (
    "You are given examples from a language you have never seen. "
    "Extract only PROVABLE OBSERVATIONS about the examples themselves. "
    "Rules:\n"
    "1. Every observation must be directly supported by specific examples.\n"
    "2. Every observation must cite the example numbers or specific words it comes from.\n"
    "3. Do NOT infer meanings, grammar, or rules beyond what the examples show.\n"
    "4. Do NOT speculate. If unsure, do not include it.\n"
    "5. Format: one observation per line, each ending with '(from: <specific examples/words>)'.\n"
    "Example format:\n"
    "  'ka' appears in examples 1, 3, 5 -- always followed by a noun (from: examples 1, 3, 5)\n"
    "  Word order in translations: subject before verb (from: examples 2, 4)\n"
    "Return 3-8 observations."
)

def arm_observation_then_answer(context, query, task_type):
    """
    Two-call arm. First call: extract observations (no query seen).
    Second call: answer with observations + context.
    """
    # Call 1: observations only, query NOT included
    obs_msgs = [
        {"role": "system", "content": OBSERVATION_SYS},
        {"role": "user",   "content": (
            f"Task type: {task_type}\n\n"
            f"EXAMPLES:\n{context.strip()}\n\n"
            f"Extract provable observations from these examples."
        )}
    ]
    observations = generate(obs_msgs, max_new_tokens=256)

    # Filter out un-cited lines (hallucination guard)
    obs_lines = [l.strip() for l in observations.splitlines() if l.strip()]
    cited = [l for l in obs_lines if "(from:" in l.lower() or "from example" in l.lower()
             or re.search(r"example[s]?\s*\d", l.lower())]
    if len(cited) < 2:
        # Observation extraction failed the grounding check
        return None, observations

    grounded_obs = "\n".join(cited)

    # Call 2: answer using observations + examples
    ans_msgs = [
        {"role": "system", "content": BASELINE_SYS},
        {"role": "user",   "content": (
            f"OBSERVATIONS ABOUT THE LANGUAGE:\n{grounded_obs}\n\n"
            f"EXAMPLES:\n{context.strip()}\n\n"
            f"{query.strip()}"
        )}
    ]
    return generate(ans_msgs, max_new_tokens=512), observations

# ---------------------------------------------------------------------------
# Answer extraction + scoring (same for all arms)
# ---------------------------------------------------------------------------
def extract_answers(text, task_type, query):
    lines = [ln.strip() for ln in text.splitlines() if ln.strip()]
    ans = [normalize(task_type, ln) for ln in lines]
    n = n_expected(query)
    if n:
        if len(ans) < n: ans = ans + [ans[-1] if ans else ""] * (n - len(ans))
        elif len(ans) > n: ans = ans[:n]
    if not ans: ans = [""]
    return ans

def score_answers(predicted, gold_raw, task_type, query):
    if isinstance(gold_raw, str):
        import ast
        try:
            gold_list = ast.literal_eval(gold_raw)
            if not isinstance(gold_list, list): gold_list = [str(gold_raw)]
        except Exception:
            gold_list = [g.strip() for g in gold_raw.split("\n") if g.strip()]
    else:
        gold_list = list(gold_raw)
    gold_norm = [normalize(task_type, str(g)) for g in gold_list]
    n = len(gold_norm)
    if len(predicted) != n:
        if len(predicted) < n:
            predicted = predicted + [predicted[-1] if predicted else ""] * (n - len(predicted))
        else:
            predicted = predicted[:n]
    correct = sum(1 for p, g in zip(predicted, gold_norm)
                  if p.strip().lower() == g.strip().lower())
    return correct, n

# ---------------------------------------------------------------------------
# MAIN LOOP: three arms, deterministic, no task leakage
# ---------------------------------------------------------------------------
results = []
t0 = time.time()

print(f"Running 3-arm diagnostic on {len(df)} problems...")
print(f"{'idx':>4} {'task':>12} {'A_base':>7} {'B_ext':>7} {'C_obs':>7} {'cited':>6} {'t':>6}")
print("-" * 65)

for idx, row in df.iterrows():
    context   = str(row.get("context", "")).strip()
    query     = str(row.get("query", "")).strip()
    answer    = row.get("answer", "")
    task_type = str(row.get("task_type", "translation")).strip()

    if not context or not query:
        continue

    # ARM A: baseline
    try:
        text_a = arm_baseline(context, query)
        ans_a  = extract_answers(text_a, task_type, query)
        em_a_c, em_a_n = score_answers(ans_a, answer, task_type, query)
        em_a = em_a_c / max(em_a_n, 1)
    except Exception as e:
        print(f"  arm A error row {idx}: {e}")
        em_a = 0.0

    # ARM B: baseline extended (compute control)
    try:
        text_b = arm_baseline_extended(context, query)
        ans_b  = extract_answers(text_b, task_type, query)
        em_b_c, em_b_n = score_answers(ans_b, answer, task_type, query)
        em_b = em_b_c / max(em_b_n, 1)
    except Exception as e:
        print(f"  arm B error row {idx}: {e}")
        em_b = em_a

    # ARM C: observations then answer
    try:
        text_c, obs_raw = arm_observation_then_answer(context, query, task_type)
        if text_c is None:
            em_c = em_a  # observation grounding check failed -- fall back
            cited_ok = False
        else:
            ans_c = extract_answers(text_c, task_type, query)
            em_c_c, em_c_n = score_answers(ans_c, answer, task_type, query)
            em_c = em_c_c / max(em_c_n, 1)
            cited_ok = True
    except Exception as e:
        print(f"  arm C error row {idx}: {e}")
        em_c = em_a
        cited_ok = False

    elapsed = time.time() - t0
    print(f"{idx:>4} {task_type:>12} {em_a:>7.3f} {em_b:>7.3f} {em_c:>7.3f} {str(cited_ok):>6} {elapsed:>5.0f}s")

    results.append({
        "idx":       idx,
        "task_type": task_type,
        "em_A":      em_a,
        "em_B":      em_b,
        "em_C":      em_c,
        "delta_BA":  em_b - em_a,   # compute effect only
        "delta_CA":  em_c - em_a,   # observations vs baseline
        "delta_CB":  em_c - em_b,   # observations vs pure extra compute
        "cited_ok":  cited_ok,
    })

# ---------------------------------------------------------------------------
# ANALYSIS: three deltas tell three different stories
# ---------------------------------------------------------------------------
import pandas as _pd
rdf = _pd.DataFrame(results)

print("\n" + "=" * 65)
print("DIAGNOSTIC SUMMARY (3 arms)")
print("=" * 65)

print(f"\nTotal problems: {len(rdf)}")
print(f"Observation grounding passed: {rdf['cited_ok'].sum()} / {len(rdf)}")

print(f"\nMean EM by arm:")
print(f"  A (baseline, greedy):        {rdf['em_A'].mean():.4f}")
print(f"  B (baseline, extra tokens):  {rdf['em_B'].mean():.4f}  (delta from A: {rdf['delta_BA'].mean():+.4f})")
print(f"  C (observations + answer):   {rdf['em_C'].mean():.4f}  (delta from A: {rdf['delta_CA'].mean():+.4f})")

print(f"\nCRITICAL COMPARISON: C - B = {rdf['delta_CB'].mean():+.4f}")
print("  This isolates the MECHANISM effect from the COMPUTE effect.")
print("  If C - B is positive, observations helped BEYOND just extra tokens.")
print("  If C - B is near zero, the gain was just from more compute.")

print("\nBy task type (delta_CA: observations effect, delta_CB: mechanism effect):")
for tt, grp in rdf.groupby("task_type"):
    print(f"  {tt:>14}: A={grp['em_A'].mean():.3f} B={grp['em_B'].mean():.3f} C={grp['em_C'].mean():.3f}  "
          f"C-A={grp['delta_CA'].mean():+.3f}  C-B={grp['delta_CB'].mean():+.3f}  n={len(grp)}")

print("\n" + "=" * 65)
print("GO / NO-GO DECISION")
print("=" * 65)

delta_CA = rdf['delta_CA'].mean()  # observations vs pure baseline
delta_CB = rdf['delta_CB'].mean()  # observations vs compute-matched baseline

if delta_CB > 0.01:
    print(f"GO: C-B = {delta_CB:+.4f} > 0.01")
    print("Observations help BEYOND extra compute. The mechanism is real.")
    print("Proceed to Phase 2: consensus over two independent observation runs.")
elif delta_CA > 0.01 and delta_CB <= 0.01:
    print(f"CONFOUND: C-A = {delta_CA:+.4f} > 0 but C-B = {delta_CB:+.4f} near zero.")
    print("The 'gain' is from extra compute, not from the observation mechanism.")
    print("Do not invest in Pass 2 architecture -- the mechanism does nothing.")
    print("Consider: is there a cheaper way to give the model more compute per problem?")
elif delta_CA <= 0.01:
    print(f"NO-GO: C-A = {delta_CA:+.4f}, no improvement over baseline.")
    print("Neither observations nor extra compute helps this model on this benchmark.")
    print("The premise is false. 0.121 is the ceiling for this inference approach.")
else:
    print(f"NEGATIVE: C-B = {delta_CB:+.4f} < 0.")
    print("Observations actively HURT compared to just running the baseline with more tokens.")
    print("Do not submit anything based on this architecture.")

print(f"\nTotal diagnostic time: {time.time()-t0:.0f}s")
print(f"Time per problem: {(time.time()-t0)/max(len(rdf),1):.1f}s")