scripts/analyze_passk_diagnostics.py

#!/usr/bin/env python
"""Build pass@10 plots and aligned logit/state summaries for detected loops."""

import csv
import json
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
from transformers import AutoTokenizer


ROOT = Path("/workspace/rwkv_doomloop_experiment")
RUNS = [
    (
        "RWKV 1.5B",
        "rwkv_1_5b",
        "11_rwkv_1_5b_pass10_sampled_autostop",
        Path("/workspace/.hf_home/hub/models--RWKV--RWKV7-G1j-1.5B-20260831/snapshots/2c18b29ab7fbece25ff6112281eea0fa41fcb30f"),
        True,
    ),
    (
        "RWKV 2.9B",
        "rwkv_2_9b",
        "12_rwkv_2_9b_pass10_sampled_autostop",
        ROOT / "models" / "RWKV7-G1j-2.9B-20260831",
        True,
    ),
    (
        "OLMo 2 1B",
        "olmo2_1b",
        "13_olmo2_1b_pass10_sampled_autostop",
        ROOT / "models" / "OLMo-2-0425-1B-Instruct",
        False,
    ),
    (
        "Qwen3.5 2B Base",
        "qwen3_5_2b_base",
        "14_qwen3_5_2b_base_pass10_sampled_autostop",
        ROOT / "models" / "Qwen3.5-2B-Base",
        False,
    ),
]
OUT = ROOT / "runs" / "passk_10_sampled_autostop"


def read_json(path):
    return json.loads(path.read_text(encoding="utf-8"))


def read_csv(path):
    with path.open(encoding="utf-8-sig", newline="") as stream:
        return list(csv.DictReader(stream))


def outcome_color(reason):
    if reason == "doom_loop_detected":
        return "#c44e52"
    if reason == "context_window_limit":
        return "#dd9a27"
    if reason == "max_new_tokens_reached":
        return "#8172b2"
    if reason == "eos":
        return "#4c72b0"
    return "#777777"


def draw_sample_lengths():
    fig, axes = plt.subplots(len(RUNS), 1, figsize=(12, 3.1 * len(RUNS)), sharex=True)
    for ax, (label, _key, run_dir, _tokenizer_path, _is_rwkv) in zip(axes, RUNS):
        base = ROOT / "runs" / run_dir
        rows = [read_json(base / f"sample_{i:02d}" / "summary.json") for i in range(1, 11)]
        lengths = [row["generated_tokens"] for row in rows]
        colors = [outcome_color(row["terminal_reason"]) for row in rows]
        ax.bar(range(1, 11), lengths, color=colors, width=0.72)
        ax.set_ylabel("response tokens")
        ax.set_title(
            f"{label}: correct {sum(row['automatic_answer_match'] for row in rows)}/10; "
            f"mean {np.mean(lengths):.1f} tokens",
            loc="left",
            fontsize=10,
        )
        ax.grid(axis="y", alpha=0.22)
        for x, y in enumerate(lengths, start=1):
            ax.text(x, y + max(12, max(lengths) * 0.012), str(y), ha="center", fontsize=8)
    axes[-1].set_xticks(range(1, 11), [str(i) for i in range(1, 11)])
    axes[-1].set_xlabel("sample index (seed 260927 + index)")
    fig.suptitle("Sampled pass@10: output length and terminal condition", y=0.995, fontsize=14)
    fig.legend(
        handles=[
            plt.Rectangle((0, 0), 1, 1, color=color, label=label)
            for label, color in [
                ("EOS", "#4c72b0"),
                ("doom-loop stop", "#c44e52"),
                ("context cap", "#dd9a27"),
                ("4096-token cap", "#8172b2"),
            ]
        ],
        loc="upper right",
        ncol=2,
        fontsize=9,
    )
    fig.tight_layout(rect=(0, 0, 1, 0.96))
    fig.savefig(OUT / "sample_lengths_and_stops.png", dpi=180)
    plt.close(fig)


def loop_comparison(label, model_key, run_dir, sample_idx, tokenizer):
    sample_dir = ROOT / "runs" / run_dir / f"sample_{sample_idx:02d}"
    summary = read_json(sample_dir / "summary.json")
    event = summary["doom_loop_event"]
    rows = read_csv(sample_dir / "token_metrics.csv")
    by_step = {int(row["predicts_response_token_step"]): row for row in rows}
    first = event["first_start_step"]
    repeat = event["repeat_start_step"]
    span = event["repeated_span_tokens"]
    first_rows = [by_step[first + offset] for offset in range(span)]
    repeat_rows = [by_step[repeat + offset] for offset in range(span)]
    token_ids_match = all(
        int(left["generated_token_id"]) == int(right["generated_token_id"])
        for left, right in zip(first_rows, repeat_rows)
    )
    generated_ids = [int(row["generated_token_id"]) for row in rows]
    span_text = tokenizer.decode(
        generated_ids[first - 1 : first - 1 + span],
        skip_special_tokens=True,
        clean_up_tokenization_spaces=False,
    )

    fields = [
        "raw_top1_probability",
        "raw_entropy_nats",
        "raw_top1_top2_margin",
        "sampled_top1_probability",
        "sampled_entropy_nats",
    ]
    if "wkv_rms_mean" in first_rows[0]:
        fields += [
            "wkv_rms_mean",
            "att_shift_rms_mean",
            "ffn_shift_rms_mean",
            "wkv_delta_rms",
            "att_shift_delta_rms",
            "ffn_shift_delta_rms",
        ]
    stats = {}
    for field in fields:
        a = np.array([float(row[field]) for row in first_rows])
        b = np.array([float(row[field]) for row in repeat_rows])
        stats[field] = {
            "first_span_mean": float(a.mean()),
            "repeat_span_mean": float(b.mean()),
            "repeat_minus_first_mean": float((b - a).mean()),
            "first_span_min": float(a.min()),
            "first_span_max": float(a.max()),
            "repeat_span_min": float(b.min()),
            "repeat_span_max": float(b.max()),
        }

    result = {
        "model_key": model_key,
        "model_label": label,
        "sample_index": sample_idx,
        "detector_event": event,
        "token_ids_match_exactly": token_ids_match,
        "repeated_span_text": span_text,
        "rolling_dup8_rate_at_detection": event["rolling_dup8_rate_64"],
        "aligned_metrics": stats,
    }
    (sample_dir / "loop_span_comparison.json").write_text(
        json.dumps(result, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
    )

    state_fields = [field for field in fields if "rms" in field]
    plot_fields = [
        "raw_top1_probability",
        "raw_entropy_nats",
    ]
    if state_fields:
        plot_fields += ["wkv_rms_mean", "att_shift_rms_mean"]
    fig, axes = plt.subplots(len(plot_fields), 1, figsize=(10, 2.3 * len(plot_fields)), sharex=True)
    if len(plot_fields) == 1:
        axes = [axes]
    x = np.arange(1, span + 1)
    for ax, field in zip(axes, plot_fields):
        a = [float(row[field]) for row in first_rows]
        b = [float(row[field]) for row in repeat_rows]
        ax.plot(x, a, label=f"first occurrence ({first}–{first + span - 1})", linewidth=1.3)
        ax.plot(x, b, label=f"repeat ({repeat}–{repeat + span - 1})", linewidth=1.3)
        ax.set_ylabel(field.replace("_", "\n"), fontsize=8)
        ax.grid(alpha=0.22)
        ax.legend(loc="best", fontsize=8)
    axes[-1].set_xlabel("token offset within repeated span")
    fig.suptitle(
        f"{label} sample {sample_idx}: logits/state at an exact {span}-token repeat",
        y=1.0,
        fontsize=12,
    )
    fig.tight_layout()
    fig.savefig(sample_dir / "loop_span_logits_state.png", dpi=180, bbox_inches="tight")
    plt.close(fig)
    return result


def main():
    OUT.mkdir(parents=True, exist_ok=True)
    draw_sample_lengths()
    loop_results = []
    for label, model_key, run_dir, tokenizer_path, is_rwkv in RUNS:
        base = ROOT / "runs" / run_dir
        tokenizer = None
        for idx in range(1, 11):
            summary = read_json(base / f"sample_{idx:02d}" / "summary.json")
            if summary["doom_loop_event"]:
                if tokenizer is None:
                    tokenizer = AutoTokenizer.from_pretrained(
                        tokenizer_path,
                        trust_remote_code=is_rwkv,
                        local_files_only=True,
                    )
                loop_results.append(
                    loop_comparison(label, model_key, run_dir, idx, tokenizer)
                )
    (OUT / "loop_span_comparisons.json").write_text(
        json.dumps(loop_results, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
    )
    print(
        json.dumps(
            {
                "sample_plot": str(OUT / "sample_lengths_and_stops.png"),
                "loop_events_analyzed": len(loop_results),
                "loop_samples": [
                    f"{item['model_key']}/sample_{item['sample_index']:02d}"
                    for item in loop_results
                ],
            },
            ensure_ascii=False,
        )
    )


if __name__ == "__main__":
    main()