external.py

"""Stage 7: independent benchmark content, two tracks.

Both environments so far were built here, which is grading my own homework.
TruthfulQA is independent and was picked in EVALUATION.md because it is built to
bait confident false answers. Each item reduces to a two choice task with
objective ground truth (Best Answer vs Best Incorrect Answer).

Track A runs the benchmark UNMODIFIED: one shot, no repeats, no feedback, no
shift. The gate cannot be computed there and is reported as N/A rather than
faked -- with no repeated items and no reward there is no belief to go stale.
Track A tests only the premise of claim 1 on independent data.

Track B is TruthfulQA-stream: real questions and real answer pairs presented as a
repeated item stream with reward feedback, where the rewarded answer flips for
some items at the shift step. This is NOT TruthfulQA any more and is labelled
modified everywhere it appears.

Track B emits the same CSV columns as run.py, so every existing scorer runs
against it unchanged.
"""

import argparse
import csv
import json
import os
import random
from collections import Counter
from concurrent.futures import ThreadPoolExecutor

import pandas as pd

from model import chat_many

SYSTEM = "Answer with exactly A or B. No other words."


def load_items(path, n, seed=0):
    """Return items as {q, true, false}. TruthfulQA ships its own best incorrect
    answer; HotpotQA does not, so a distractor is built (see _hotpot)."""
    if path.endswith(".parquet"):
        return _hotpot(path, n, seed, yesno_only=os.environ.get("HOTPOT_YESNO") == "1")
    if path.endswith(".jsonl"):
        return _treecut(path, n, seed)
    df = pd.read_csv(path).dropna(subset=["Question", "Best Answer", "Best Incorrect Answer"])
    df = df.sample(n=min(n, len(df)), random_state=seed)
    return [{"q": r["Question"].strip(),
             "true": str(r["Best Answer"]).strip(),
             "false": str(r["Best Incorrect Answer"]).strip()}
            for _, r in df.iterrows()]


def _treecut(path, n, seed=0):
    """TreeCut math word problems. Numeric answers, so the distractor is another
    item's answer, which keeps both options the same shape.

    NOTE: this uses TreeCut's ANSWERABLE half only. The unanswerable half is the
    dataset's actual point and is a different experiment (does the detector flag a
    confident answer to a question with no answer), recorded as future work rather
    than silently folded in here."""
    rows = [json.loads(l) for l in open(path)]
    rng = random.Random(seed)
    rng.shuffle(rows)
    answers = [str(r["answer"]).strip() for r in rows]
    items = []
    for r in rows:
        ans = str(r["answer"]).strip()
        near = [a for a in answers if a != ans]
        if not near:
            continue
        items.append({"q": str(r["problem"]).strip(),
                      "true": ans, "false": rng.choice(near)})
        if len(items) >= n:
            break
    return items


def _hotpot(path, n, seed=0, yesno_only=False):
    """HotpotQA has no wrong answer supplied, so one is constructed.

    yes/no questions get the opposite word. Everything else gets the answer from a
    DIFFERENT question of similar length, so the distractor is a plausible entity of
    the right shape rather than something ruled out by format alone. Questions are
    asked without their supporting context, which is the point: the benchmark is
    being used as a source of items the model answers from memory."""
    df = pd.read_parquet(path)
    df = df[df["answer"].str.len() < 30].dropna(subset=["question", "answer"])
    # drop truncated or malformed answers; a distractor that is obviously broken text
    # would be rejected on format alone and would not test anything
    ok = df["answer"].str.match(r"^[A-Za-z0-9][^\"]*[A-Za-z0-9.)]$") & (df["answer"].str.len() > 2)
    df = df[ok]
    if yesno_only:
        # constructed distractors are trivially rejectable for a strong model, which
        # pinned accuracy at 0.997 and left one class empty. yes/no items are a genuine
        # forced choice: both options are well formed and neither can be ruled out on shape.
        df = df[df["answer"].str.lower().isin(["yes", "no"])]
    df = df.sample(n=min(n * 3, len(df)), random_state=seed).reset_index(drop=True)
    rng = random.Random(seed)
    pool = [a for a in df["answer"].tolist() if a.lower() not in ("yes", "no")]
    items = []
    for _, r in df.iterrows():
        ans = str(r["answer"]).strip()
        if ans.lower() in ("yes", "no"):
            wrong = "no" if ans.lower() == "yes" else "yes"
        else:
            near = [a for a in pool if a != ans and abs(len(a) - len(ans)) <= 6]
            if not near:
                continue
            wrong = rng.choice(near)
        items.append({"q": str(r["question"]).strip(), "true": ans, "false": wrong})
        if len(items) >= n:
            break
    return items


def ask(item, k, flip, history_note=""):
    """k paraphrases. `flip` decides which option letter carries the true answer,
    so the letter itself never leaks the answer."""
    a, b = (item["true"], item["false"]) if not flip else (item["false"], item["true"])
    stems = [
        f"Question: {item['q']}\nA) {a}\nB) {b}\nWhich is correct?",
        f"{item['q']}\nOption A: {a}\nOption B: {b}\nAnswer A or B.",
        f"Consider: {item['q']}\nA. {a}\nB. {b}\nPick one.",
        f"For the question '{item['q']}', is the right answer A) {a} or B) {b}?",
        f"{item['q']}\n(A) {a}\n(B) {b}\nRespond with the correct letter.",
    ]
    return [stems[i % len(stems)] + history_note for i in range(k)]


def parse_ab(text):
    t = (text or "").strip().upper()
    for ch in t:
        if ch in "AB":
            return ch
    return "A"


# ---------------------------------------------------------------- track A

def track_a(items, k, model, out, workers=8):
    """Unmodified benchmark: one shot per question, no feedback, no shift."""
    rng = random.Random(0)
    jobs = [(i, it, rng.random() < 0.5) for i, it in enumerate(items)]

    def one(job):
        i, it, flip = job
        texts = chat_many(ask(it, k, flip), model, max_tokens=4, system=SYSTEM)
        picks = [parse_ab(t) for t in texts]
        chosen = Counter(picks).most_common(1)[0][0]
        truth = "B" if flip else "A"
        return {
            "step": i, "phase": "static", "cue": f"q{i}",
            "chosen": chosen, "correct": truth,
            "reward": int(chosen == truth),
            "agreement": round(picks.count(chosen) / len(picks), 3),
            "branches": "|".join(picks),
        }

    with ThreadPoolExecutor(max_workers=workers) as pool:
        rows = list(pool.map(one, jobs))
    write(out, rows)
    return rows


# ---------------------------------------------------------------- track B

def track_b(items, seed, steps, shift_at, k, model, out, changed_frac=0.5):
    """TruthfulQA-stream (MODIFIED): repeated items, reward feedback, rule shift.

    Pre shift the truthful answer is rewarded. At the shift step a subset of items
    flips, so the previously wrong answer becomes the rewarded one. The agent is
    told nothing; it only ever sees reward."""
    rng = random.Random(seed)
    pool_items = list(items)
    flipped = set(rng.sample(range(len(pool_items)),
                             max(1, round(changed_frac * len(pool_items)))))
    # which letter the true answer sits behind, fixed per item so it is not a cue
    letter_flip = {i: rng.random() < 0.5 for i in range(len(pool_items))}

    rows, history = [], []
    for step in range(steps):
        idx = rng.randrange(len(pool_items))
        it = pool_items[idx]
        post = step >= shift_at
        # pre shift: truthful is rewarded. post shift: flipped items reward the other one.
        truth_is_true_answer = not (post and idx in flipped)
        lf = letter_flip[idx]
        # letter carrying the currently rewarded answer
        if truth_is_true_answer:
            correct = "B" if lf else "A"
        else:
            correct = "A" if lf else "B"

        seen = [h for h in history if h["cue"] == f"q{idx}"][-4:]
        note = ("\nEarlier attempts on this question: "
                + "; ".join(f"{h['chosen']} was {'right' if int(h['reward']) else 'wrong'}"
                            for h in seen)) if seen else ""
        picks = [parse_ab(t) for t in chat_many(
            ask(it, k, lf, note), model, max_tokens=4, system=SYSTEM)]
        chosen = Counter(picks).most_common(1)[0][0]
        row = {
            "step": step,
            "phase": "post" if post else "pre",
            "cue": f"q{idx}",
            "chosen": chosen,
            "correct": correct,
            "reward": int(chosen == correct),
            "agreement": round(picks.count(chosen) / len(picks), 3),
            "branches": "|".join(picks),
        }
        rows.append(row)
        history.append(row)
    write(out, rows)
    return rows


def write(path, rows):
    folder = os.path.dirname(path)
    if folder:
        os.makedirs(folder, exist_ok=True)
    with open(path, "w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=list(rows[0].keys()))
        w.writeheader()
        w.writerows(rows)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--track", choices=["a", "b"], required=True)
    ap.add_argument("--data", default="data/TruthfulQA.csv")
    ap.add_argument("--items", type=int, default=200, help="track a: questions; track b: pool size")
    ap.add_argument("--seeds", type=int, default=30, help="track b only: episodes")
    ap.add_argument("--steps", type=int, default=80)
    ap.add_argument("--shift_at", type=int, default=60)
    ap.add_argument("--k", type=int, default=5)
    ap.add_argument("--changed_frac", type=float, default=0.5)
    ap.add_argument("--model", default="meta-llama/llama-3.3-70b-instruct")
    ap.add_argument("--out_dir", default="traces/stage7_truthfulqa")
    ap.add_argument("--workers", type=int, default=8)
    args = ap.parse_args()

    if args.track == "a":
        items = load_items(args.data, args.items)
        rows = track_a(items, args.k, args.model,
                       os.path.join(args.out_dir, "trackA_static.csv"), args.workers)
        acc = sum(r["reward"] for r in rows) / len(rows)
        print(f"track A: {len(rows)} questions, accuracy {acc:.3f}")
        return

    # track b: a small item pool per episode so items actually repeat
    jobs = []
    for s in range(args.seeds):
        items = load_items(args.data, 12, seed=s)
        jobs.append((s, items, os.path.join(args.out_dir, f"shift_seed{s}.csv")))

    def one(job):
        s, items, out = job
        track_b(items, s, args.steps, args.shift_at, args.k, args.model, out,
                changed_frac=args.changed_frac)
        return out

    with ThreadPoolExecutor(max_workers=args.workers) as pool:
        for i, done in enumerate(pool.map(one, jobs), 1):
            print(f"[{i}/{len(jobs)}] {done}", flush=True)


if __name__ == "__main__":
    main()