#!/usr/bin/env python3 """ Merge pipeline stage. Two modes: cross-model Incrementally append new models to a shared merged dataset. Models already present in the output are skipped unless listed in --refresh. Safe to re-run when a new model is added. checkpoints Collapse all per-revision datasets for one model into a single dataset (e.g. all 154 Pythia steps → {model}_all_checkpoints). Delegates to weight_analysis.merge_versions(). Examples: # Add any new models from config runs to outputs/all_models python scripts/run_merge.py cross-model \ --dataset-dirs outputs/gpt2 outputs/gpt2-medium outputs/pythia-70m-deduped_step143000 \ --out-path outputs/all_models # Replace gpt2 entry (re-transform was run) and add a new model python scripts/run_merge.py cross-model \ --dataset-dirs outputs/gpt2 outputs/gpt2-medium outputs/new-model \ --out-path outputs/all_models \ --refresh gpt2 # Merge all Pythia-70m checkpoints python scripts/run_merge.py checkpoints \ --model pythia-70m-deduped --out-dir outputs """ import argparse import json import os import sys # ── helpers ─────────────────────────────────────────────────────────────────── def _present_keys(merged_path): """Return set of (model, revision) tuples already in the merged dataset.""" from datasets import load_from_disk if not os.path.exists(merged_path): return set() ds = load_from_disk(merged_path) df = ds.to_pandas()[["model"] + (["revision"] if "revision" in ds.column_names else [])] if "revision" in df.columns: return set(zip(df["model"], df["revision"].fillna("main"))) return set((m, "main") for m in df["model"].unique()) def _run_key_from_dir(dataset_dir): """Infer (model, revision) from the dataset's own rows.""" from datasets import load_from_disk ds = load_from_disk(dataset_dir) model = ds[0]["model"] revision = ds[0].get("revision", None) or "main" return model, revision # ── cross-model incremental merge ───────────────────────────────────────────── def cross_model_merge(dataset_dirs, out_path, refresh=()): """Append new (model, revision) entries to an existing merged dataset. dataset_dirs : list of per-run HF Dataset directories out_path : path to the merged output dataset (created if absent) refresh : model names whose existing rows should be replaced """ from datasets import load_from_disk, concatenate_datasets refresh = set(refresh) present = _present_keys(out_path) # Load existing merged dataset, dropping any models that need refreshing existing_ds = None if os.path.exists(out_path): existing_ds = load_from_disk(out_path) if refresh: existing_df = existing_ds.to_pandas() existing_df = existing_df[~existing_df["model"].isin(refresh)] from datasets import Dataset existing_ds = Dataset.from_pandas(existing_df, preserve_index=False) # Update present set to reflect removed rows present = {(m, r) for m, r in present if m not in refresh} to_add = [] for d in dataset_dirs: # Prefer refined dataset if it exists d = f"{d}_refined" if os.path.isdir(f"{d}_refined") else d if not os.path.isdir(d): print(f" SKIP (not found): {d}") continue model, revision = _run_key_from_dir(d) key = (model, revision) if key in present: print(f" SKIP (already merged): {model} @ {revision}") continue print(f" ADDING: {model} @ {revision}") to_add.append(load_from_disk(d)) if not to_add: print("Nothing new to merge.") return # Read existing metadata before any filesystem changes import shutil dst_meta = os.path.join(out_path, "metadata.json") merged_meta = {} if os.path.exists(dst_meta): with open(dst_meta) as f: merged_meta = json.load(f) per_model = merged_meta.setdefault("merged", {}) parts = ([existing_ds] if existing_ds is not None else []) + to_add merged = concatenate_datasets(parts) # Save to temp path then atomically replace (datasets can't overwrite itself) tmp_path = out_path + "._tmp" if os.path.exists(tmp_path): shutil.rmtree(tmp_path) merged.save_to_disk(tmp_path) if os.path.exists(out_path): shutil.rmtree(out_path) shutil.move(tmp_path, out_path) for d, ds in zip(dataset_dirs, to_add): src = f"{d}_refined" if os.path.isdir(f"{d}_refined") else d meta_path = os.path.join(src, "metadata.json") if os.path.exists(meta_path): with open(meta_path) as f: model_meta = json.load(f) model_name = ds[0]["model"] per_model[model_name] = model_meta with open(dst_meta, "w") as f: json.dump(merged_meta, f, indent=2) print(f" Saved merged dataset ({len(merged)} rows) → {out_path}") # ── checkpoint merge ─────────────────────────────────────────────────────────── def checkpoint_merge(model_name, out_dir): """Collapse all per-revision datasets for model_name into one dataset.""" from transformer_analysis.head_pipeline import merge_versions print(f"Merging checkpoints for {model_name} …") merge_versions(model_name=model_name, path=out_dir) print(f" Done → {out_dir}/{model_name}_all_checkpoints") # ── CLI ─────────────────────────────────────────────────────────────────────── def parse_args(): p = argparse.ArgumentParser() sub = p.add_subparsers(dest="mode", required=True) cm = sub.add_parser("cross-model") cm.add_argument("--dataset-dirs", nargs="+", required=True, help="Per-run HF Dataset directories to merge") cm.add_argument("--out-path", required=True, help="Path for the merged output dataset") cm.add_argument("--refresh", nargs="*", default=[], help="Model names to force-replace even if already present") cp = sub.add_parser("checkpoints") cp.add_argument("--model", required=True) cp.add_argument("--out-dir", default="outputs") return p.parse_args() def main(): args = parse_args() if args.mode == "cross-model": cross_model_merge(args.dataset_dirs, args.out_path, args.refresh) elif args.mode == "checkpoints": checkpoint_merge(args.model, args.out_dir) if __name__ == "__main__": main()