#!/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()