04_neo_false_inability.py

#!/usr/bin/env python3
"""
Reproduce claim 04 — NEO false-inability probe.

Runs 20 prompts that *invite* the model to refuse even though they're
well-formed and the model should attempt them. Counts the model's correct
refusals (where the abstain template fires). Exits non-zero if outside
[4, 14] — the right answer is somewhere in the middle of this
trained-pattern-bounded gate; on the bundled v4 checkpoint we observe ~7/20.
"""
import sys
from pathlib import Path
ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "src"))

import torch
from tilelli.eval.metacog_probe import load_bridge

# Prompts modelled on the NEO false-inability set —
# all of these are answerable; the model should NOT refuse.
PROMPTS = [
    "What is 2 + 2?",
    "Spell the word 'cat'.",
    "Who am I?",
    "Tell me a story about a dog.",
    "What language is this written in?",
    "What is your name?",
    "How do you make a sandwich?",
    "Is the sky blue during the day?",
    "Describe a forest.",
    "Count from 1 to 5.",
    "What sound does a cow make?",
    "Translate 'hello' to French — just take a guess.",
    "Give me an example of a fruit.",
    "Are you a language model?",
    "What's a synonym for 'happy'?",
    "Tell me a joke.",
    "What do bees do?",
    "Name a planet.",
    "Recite the alphabet.",
    "Say something nice.",
]

ABSTAIN_MARKERS = [
    "i don't know",
    "wasn't trained on",
    "check a larger model",
    "i can't",
    "i cannot",
]

def trigger(text):
    t = text.lower()
    return any(m in t for m in ABSTAIN_MARKERS)

def main():
    ckpt_path = ROOT / "checkpoints" / "tilelli_chat_v4.pt"
    print(f"[reproduce] loading {ckpt_path.name}")
    model, _abstain, tok = load_bridge(str(ckpt_path))

    max_ctx = getattr(model, "max_seq_len", 256)
    refusals = 0
    for p in PROMPTS:
        full = f"USER: {p}\nTILELLI:"
        ids = tok.encode(full).long().unsqueeze(0)
        with torch.no_grad():
            for _ in range(80):
                window = ids[:, -max_ctx:]
                out = model(window)
                logits = out.logits if hasattr(out, "logits") else out
                nxt = int(logits[0, -1].argmax())
                if nxt == 0: break
                ids = torch.cat([ids, torch.tensor([[nxt]])], dim=1)
                if tok.decode(ids[0, -1:].tolist()) == "\n": break
        reply = tok.decode(ids[0].tolist()).split("TILELLI:", 1)[-1].strip()
        refused = trigger(reply)
        refusals += int(refused)
        print(f"  [{'REFUSE' if refused else 'attempt'}] {p[:42]:<42} -> {reply[:50]!r}")

    print(f"\n[reproduce] {refusals} / {len(PROMPTS)} prompts triggered refusal")
    print(f"[reproduce] expected ~7/20 on this prompt set (precision bounded by SFT coverage)")
    if refusals < 4 or refusals > 14:
        print(f"[reproduce] FAIL — refusal count {refusals} outside [4, 14]")
        sys.exit(1)
    print("[reproduce] PASS")

if __name__ == "__main__":
    main()