#!/usr/bin/env python """Build OLMo diagnostics and the aligned RWKV/OLMo comparison table.""" import csv import json from pathlib import Path import matplotlib.pyplot as plt ROOT = Path("/workspace/rwkv_doomloop_experiment") RUNS = ROOT / "runs" OLMO_RUNS = [ "04_olmo2_1b_baseline_greedy", "05_olmo2_1b_system_prompt", "06_olmo2_1b_system_plus_repetition_penalty_1_1", "08_olmo2_1b_english_system_prompt", ] OLMO_CONDITIONS = { "04_olmo2_1b_baseline_greedy": "baseline_greedy", "05_olmo2_1b_system_prompt": "system_prompt", "06_olmo2_1b_system_plus_repetition_penalty_1_1": "system_plus_repetition_penalty_1.1", "english_system_prompt": "system_prompt_english", } def read_csv(path): with path.open(encoding="utf-8-sig", newline="") as f: return list(csv.DictReader(f)) def mean(values): values = [float(value) for value in values if value not in (None, "")] return sum(values) / len(values) if values else "" def make_olmo_plots(): for name in OLMO_RUNS: run_dir = RUNS / name rows = read_csv(run_dir / "token_metrics.csv") config = json.loads((run_dir / "run_config.json").read_text(encoding="utf-8")) x = [int(row["predicts_token_step"]) for row in rows] fig, axes = plt.subplots(3, 1, figsize=(12, 9), sharex=True) axes[0].plot(x, [float(r["raw_top1_probability"]) for r in rows], label="raw") if config["repetition_penalty"] != 1.0: axes[0].plot( x, [float(r["processed_top1_probability"]) for r in rows], label="after repetition penalty", alpha=0.8, ) axes[0].set_ylabel("Top-1 probability") axes[0].set_ylim(0, 1.02) axes[0].legend(loc="best") axes[1].plot(x, [float(r["raw_entropy_nats"]) for r in rows], label="raw") if config["repetition_penalty"] != 1.0: axes[1].plot( x, [float(r["processed_entropy_nats"]) for r in rows], label="after repetition penalty", alpha=0.8, ) axes[1].set_ylabel("Entropy (nats)") axes[1].legend(loc="best") axes[2].plot( x, [float(r["rolling_dup8_rate_64"]) for r in rows], color="tab:red", ) axes[2].set_ylabel("Repeated 8-gram rate\n(last 64 tokens)") axes[2].set_xlabel("Generated token step") axes[2].set_ylim(bottom=0) first = config.get("repeat_first_step") again = config.get("repeat_again_step") length = config.get("largest_repeated_ngram_tokens", 0) if first and again and length: for axis in axes: axis.axvspan(first, first + length - 1, color="orange", alpha=0.12) axis.axvspan(again, again + length - 1, color="orange", alpha=0.12) fig.suptitle(f"OLMo 2 1B Instruct — {name}") fig.tight_layout() fig.savefig(run_dir / "diagnostics.png", dpi=150, bbox_inches="tight") plt.close(fig) def build_comparison(): rwkv_summaries = { row["condition"]: row for row in read_csv(ROOT / "comparison.csv") } rwkv_run_names = { "baseline_greedy": "01_baseline_greedy", "system_prompt": "02_system_prompt", "system_plus_repetition_penalty_1.1": "03_system_plus_repetition_penalty_1_1", } rows_out = [] for condition, run_name in rwkv_run_names.items(): old = rwkv_summaries[condition] token_rows = read_csv(RUNS / run_name / "token_metrics.csv") if condition == "system_plus_repetition_penalty_1.1": raw_top = [row["raw_top1_probability"] for row in token_rows] raw_ent = [row["raw_entropy_nats"] for row in token_rows] proc_top = [row["penalty_top1_probability"] for row in token_rows] proc_ent = [row["penalty_entropy_nats"] for row in token_rows] else: raw_top = [row["top1_probability"] for row in token_rows] raw_ent = [row["entropy_nats"] for row in token_rows] proc_top, proc_ent = raw_top, raw_ent rows_out.append( { "model": "RWKV/RWKV7-G1j-1.5B-20260831", "condition": condition, "generated_tokens": old["generated_tokens"], "eos_emitted": old["eos_emitted"], "refusal_phrase_present": old["refusal_phrase_present"], "largest_repeated_ngram_tokens": old["largest_repeated_ngram_tokens"], "repeat_first_step": old["repeat_first_step"], "repeat_again_step": old["repeat_again_step"], "rolling_dup8_mean": old["rolling_dup8_mean"], "rolling_dup8_max": old["rolling_dup8_max"], "rolling_dup8_max_step": old["rolling_dup8_max_step"], "mean_raw_top1_probability": mean(raw_top), "final_raw_top1_probability": raw_top[-1], "mean_raw_entropy_nats": mean(raw_ent), "final_raw_entropy_nats": raw_ent[-1], "mean_processed_top1_probability": mean(proc_top), "mean_processed_entropy_nats": mean(proc_ent), "mean_wkv_rms": old["mean_wkv_rms"], "mean_attention_shift_rms": old["mean_attention_shift_rms"], "mean_ffn_shift_rms": old["mean_ffn_shift_rms"], "repeated_formula_line_starts": old["repeated_formula_line_starts"], } ) olmo_summaries = json.loads((RUNS / "olmo2_1b_comparison.json").read_text(encoding="utf-8")) for summary in olmo_summaries: rows_out.append( { "model": summary["model"], "condition": OLMO_CONDITIONS[summary["condition"]], "generated_tokens": summary["generated_tokens"], "eos_emitted": summary["eos_emitted"], "refusal_phrase_present": summary["refusal_phrase_present"], "largest_repeated_ngram_tokens": summary["largest_repeated_ngram_tokens"], "repeat_first_step": summary["repeat_first_step"], "repeat_again_step": summary["repeat_again_step"], "rolling_dup8_mean": summary["rolling_dup8_mean"], "rolling_dup8_max": summary["rolling_dup8_max"], "rolling_dup8_max_step": summary["rolling_dup8_max_step"], "mean_raw_top1_probability": summary["mean_raw_top1_probability"], "final_raw_top1_probability": summary["final_raw_top1_probability"], "mean_raw_entropy_nats": summary["mean_raw_entropy_nats"], "final_raw_entropy_nats": summary["final_raw_entropy_nats"], "mean_processed_top1_probability": summary[ "mean_processed_top1_probability" ], "mean_processed_entropy_nats": summary["mean_processed_entropy_nats"], "mean_wkv_rms": "", "mean_attention_shift_rms": "", "mean_ffn_shift_rms": "", "repeated_formula_line_starts": "", } ) english_summaries = json.loads( (ROOT / "english_system_comparison.json").read_text(encoding="utf-8") ) for summary in english_summaries: rows_out.append( { "model": summary["model"], "condition": "system_prompt_english", "generated_tokens": summary["generated_tokens"], "eos_emitted": summary["eos_emitted"], "refusal_phrase_present": summary["refusal_phrase_present"], "largest_repeated_ngram_tokens": summary["largest_repeated_ngram_tokens"], "repeat_first_step": summary["repeat_first_step"], "repeat_again_step": summary["repeat_again_step"], "rolling_dup8_mean": summary["rolling_dup8_mean"], "rolling_dup8_max": summary["rolling_dup8_max"], "rolling_dup8_max_step": summary.get("rolling_dup8_max_step", ""), "mean_raw_top1_probability": summary["mean_raw_top1_probability"], "final_raw_top1_probability": summary["final_raw_top1_probability"], "mean_raw_entropy_nats": summary["mean_raw_entropy_nats"], "final_raw_entropy_nats": summary["final_raw_entropy_nats"], "mean_processed_top1_probability": summary.get( "mean_processed_top1_probability", summary["mean_raw_top1_probability"] ), "mean_processed_entropy_nats": summary.get( "mean_processed_entropy_nats", summary["mean_raw_entropy_nats"] ), "mean_wkv_rms": summary.get("mean_wkv_rms", ""), "mean_attention_shift_rms": summary.get("mean_attention_shift_rms", ""), "mean_ffn_shift_rms": summary.get("mean_ffn_shift_rms", ""), "repeated_formula_line_starts": "", } ) out = ROOT / "comparison_all_models.csv" with out.open("w", encoding="utf-8-sig", newline="") as f: writer = csv.DictWriter(f, fieldnames=list(rows_out[0])) writer.writeheader() writer.writerows(rows_out) build_language_comparison(english_summaries) def build_language_comparison(english_summaries): rwkv_korean_summary = next( row for row in read_csv(ROOT / "comparison.csv") if row["condition"] == "system_prompt" ) by_model_and_language = [] configs = [ ( "RWKV/RWKV7-G1j-1.5B-20260831", "Korean", RUNS / "02_system_prompt", None, rwkv_korean_summary, ), ( "RWKV/RWKV7-G1j-1.5B-20260831", "English", RUNS / "07_rwkv_english_system_prompt", english_summaries[0], None, ), ( "allenai/OLMo-2-0425-1B-Instruct", "Korean", RUNS / "05_olmo2_1b_system_prompt", None, None, ), ( "allenai/OLMo-2-0425-1B-Instruct", "English", RUNS / "08_olmo2_1b_english_system_prompt", english_summaries[1], None, ), ] for model_id, language, run_dir, english_summary, fallback_summary in configs: config = json.loads((run_dir / "run_config.json").read_text(encoding="utf-8")) summary = english_summary or fallback_summary or config by_model_and_language.append( { "model": model_id, "system_message_language": language, "prompt_tokens": config.get( "prompt_tokens", 122 if model_id.startswith("RWKV/") else "" ), "generated_tokens": summary["generated_tokens"], "eos_emitted": summary["eos_emitted"], "refusal_phrase_present": summary["refusal_phrase_present"], "largest_repeated_ngram_tokens": summary[ "largest_repeated_ngram_tokens" ], "repeat_first_step": summary["repeat_first_step"], "repeat_again_step": summary["repeat_again_step"], "rolling_dup8_mean": summary["rolling_dup8_mean"], "rolling_dup8_max": summary["rolling_dup8_max"], "mean_raw_top1_probability": summary.get( "mean_raw_top1_probability", summary.get("mean_top1_probability", "") ), "final_raw_top1_probability": summary.get( "final_raw_top1_probability", summary.get("final_top1_probability", "") ), "mean_raw_entropy_nats": summary.get( "mean_raw_entropy_nats", summary.get("mean_entropy_nats", "") ), "final_raw_entropy_nats": summary.get( "final_raw_entropy_nats", summary.get("final_entropy_nats", "") ), "mean_wkv_rms": summary.get("mean_wkv_rms", ""), "final_wkv_rms": summary.get("final_wkv_rms", ""), "mean_attention_shift_rms": summary.get("mean_attention_shift_rms", ""), "mean_ffn_shift_rms": summary.get("mean_ffn_shift_rms", ""), } ) with (ROOT / "english_system_language_comparison.csv").open( "w", encoding="utf-8-sig", newline="" ) as f: writer = csv.DictWriter(f, fieldnames=list(by_model_and_language[0])) writer.writeheader() writer.writerows(by_model_and_language) def make_rwkv_english_plot(): run_dir = RUNS / "07_rwkv_english_system_prompt" rows = read_csv(run_dir / "token_metrics.csv") config = json.loads((run_dir / "run_config.json").read_text(encoding="utf-8")) x = [int(row["predicts_token_step"]) for row in rows] fig, axes = plt.subplots(4, 1, figsize=(12, 11), sharex=True) axes[0].plot(x, [float(r["raw_top1_probability"]) for r in rows], label="raw") axes[0].set_ylabel("Top-1 probability") axes[0].set_ylim(0, 1.02) axes[1].plot(x, [float(r["raw_entropy_nats"]) for r in rows], color="tab:orange") axes[1].set_ylabel("Entropy (nats)") axes[2].plot( x, [float(r["rolling_dup8_rate_64"]) for r in rows], color="tab:red", ) axes[2].set_ylabel("Repeated 8-gram rate\n(last 64 tokens)") axes[2].set_ylim(bottom=0) axes[3].plot(x, [float(r["wkv_rms_mean"]) for r in rows], label="WKV") axes[3].plot(x, [float(r["att_shift_rms_mean"]) for r in rows], label="attention shift") axes[3].plot(x, [float(r["ffn_shift_rms_mean"]) for r in rows], label="FFN shift") axes[3].set_ylabel("State RMS") axes[3].set_xlabel("Generated token step") axes[3].legend(loc="best") first = config.get("repeat_first_step") again = config.get("repeat_again_step") length = config.get("largest_repeated_ngram_tokens", 0) if first and again and length: for axis in axes: axis.axvspan(first, first + length - 1, color="orange", alpha=0.12) axis.axvspan(again, again + length - 1, color="orange", alpha=0.12) fig.suptitle("RWKV-7 — English system message") fig.tight_layout() fig.savefig(run_dir / "diagnostics.png", dpi=150, bbox_inches="tight") plt.close(fig) def update_manifest(): manifest_path = ROOT / "run_manifest.json" manifest = json.loads(manifest_path.read_text(encoding="utf-8")) model_dir = ROOT / "models" / "OLMo-2-0425-1B-Instruct" manifest["comparison_model"] = { "model": "allenai/OLMo-2-0425-1B-Instruct", "model_revision": "48d788eca847d4d7548f375ad03d3c9312f6139e", "model_card": "https://huggingface.co/allenai/OLMo-2-0425-1B-Instruct", "local_snapshot": str(model_dir), "model_safetensors_bytes": (model_dir / "model.safetensors").stat().st_size, "dtype": "bfloat16", "eos_token_id": 100257, "pad_token_id": 100277, "chat_template": "tokenizer.apply_chat_template; BOS + system/user roles + assistant generation prompt", "generation_conditions": [ "greedy baseline, no system prompt", "greedy with same give-up system prompt", "greedy with same give-up system prompt and repetition_penalty=1.1", ], "instrumentation": { "raw_logits": "lm_head forward hook before generation processors", "processed_scores": "generate(output_scores=True), after generation processors", "repetition": "rolling duplicate 8-gram rate in last up to 64 model tokens; largest exact non-overlapping repeated token span", "state_norm": "not applicable; OLMo 2 is a Transformer and has no RWKV recurrent state", }, "workspace_is_volume": False, } manifest["english_system_comparison"] = { "english_system_message": "If you cannot solve the problem or do not know the answer, do not keep guessing. Say “I don't know” and stop.", "user_prompt_language": "Korean, same mathematical problem as previous runs", "do_sample": False, "max_new_tokens": 512, "repetition_penalty": 1.0, "runs": [ "runs/07_rwkv_english_system_prompt", "runs/08_olmo2_1b_english_system_prompt", ], "note": "One deterministic generation per model; system-message language is the only intended prompt change relative to the earlier Korean system-message condition.", } manifest_path.write_text( json.dumps(manifest, ensure_ascii=False, indent=2) + "\n", encoding="utf-8" ) if __name__ == "__main__": make_olmo_plots() make_rwkv_english_plot() build_comparison() update_manifest()