Add weak-layer stage2 dataset builder

scripts/build_weak_layer_dataset.py

@@ -0,0 +1,119 @@
+#!/usr/bin/env python3
+"""Build a weak-layer-focused second-stage SFT dataset.
+
+Creates local parquet files:
+  <output_dir>/train.parquet
+  <output_dir>/validation.parquet
+  <output_dir>/manifest.json
+
+The train split contains:
+- all examples from weak target layers,
+- extra duplicated rows for very rare weak layers up to a configurable minimum,
+- a replay buffer sampled from non-weak layers to reduce catastrophic forgetting.
+
+All eval/test reporting should still use the official research-sota OOD splits.
+"""
+import argparse
+import json
+from pathlib import Path
+
+import pandas as pd
+from datasets import load_dataset
+
+WEAK_LAYERS_DEFAULT = [
+    "o1_nrm",
+    "a1_policy",
+    "tmf921_lifecycle_report",
+    "tmf921_lifecycle_monitor",
+    "tmf921_lifecycle_scale",
+]
+
+
+def parse_args():
+    p = argparse.ArgumentParser()
+    p.add_argument("--dataset", default="nraptisss/TMF921-intent-to-config-research-sota")
+    p.add_argument("--train_split", default="train_sota")
+    p.add_argument("--validation_split", default="validation")
+    p.add_argument("--output_dir", required=True)
+    p.add_argument("--weak_layers", nargs="+", default=WEAK_LAYERS_DEFAULT)
+    p.add_argument("--rare_min_per_layer", type=int, default=1500, help="Duplicate weak layers with fewer rows to this minimum")
+    p.add_argument("--replay_ratio", type=float, default=0.30, help="Replay rows as fraction of final weak rows")
+    p.add_argument("--seed", type=int, default=42)
+    return p.parse_args()
+
+
+def main():
+    args = parse_args()
+    out = Path(args.output_dir)
+    out.mkdir(parents=True, exist_ok=True)
+
+    ds = load_dataset(args.dataset)
+    train = ds[args.train_split].to_pandas()
+    val = ds[args.validation_split].to_pandas()
+
+    weak_layers = set(args.weak_layers)
+    weak_parts = []
+    layer_counts_before = {}
+    layer_counts_after = {}
+
+    for layer in args.weak_layers:
+        part = train[train["target_layer"] == layer].copy()
+        layer_counts_before[layer] = int(len(part))
+        if len(part) == 0:
+            continue
+        if len(part) < args.rare_min_per_layer:
+            reps = []
+            needed = args.rare_min_per_layer - len(part)
+            for i in range(needed):
+                r = part.iloc[i % len(part)].copy(deep=True)
+                original_id = r["id"]
+                r["id"] = f"{original_id}-stage2weak-{i:05d}"
+                r["is_augmented"] = True
+                r["augmentation_type"] = f"stage2_weak_duplicate_{layer}"
+                r["source_id"] = r.get("source_id", original_id)
+                reps.append(r)
+            if reps:
+                part = pd.concat([part, pd.DataFrame(reps)], ignore_index=True)
+        layer_counts_after[layer] = int(len(part))
+        weak_parts.append(part)
+
+    weak_df = pd.concat(weak_parts, ignore_index=True) if weak_parts else pd.DataFrame(columns=train.columns)
+    nonweak = train[~train["target_layer"].isin(weak_layers)].copy()
+    replay_n = min(len(nonweak), int(len(weak_df) * args.replay_ratio))
+    replay = nonweak.sample(n=replay_n, random_state=args.seed).copy() if replay_n > 0 else pd.DataFrame(columns=train.columns)
+    replay["augmentation_type"] = replay.get("augmentation_type", "none").astype(str) + "+stage2_replay"
+
+    stage2 = pd.concat([weak_df, replay], ignore_index=True).sample(frac=1.0, random_state=args.seed).reset_index(drop=True)
+    stage2["stage2_role"] = stage2["target_layer"].apply(lambda x: "weak" if x in weak_layers else "replay")
+    val = val.copy()
+    val["stage2_role"] = "validation"
+
+    train_path = out / "train.parquet"
+    val_path = out / "validation.parquet"
+    stage2.to_parquet(train_path, index=False)
+    val.to_parquet(val_path, index=False)
+
+    manifest = {
+        "source_dataset": args.dataset,
+        "train_split": args.train_split,
+        "validation_split": args.validation_split,
+        "output_dir": str(out),
+        "weak_layers": args.weak_layers,
+        "rare_min_per_layer": args.rare_min_per_layer,
+        "replay_ratio": args.replay_ratio,
+        "seed": args.seed,
+        "rows_train_stage2": int(len(stage2)),
+        "rows_validation": int(len(val)),
+        "weak_rows_total_after_duplication": int(len(weak_df)),
+        "replay_rows": int(len(replay)),
+        "layer_counts_before": layer_counts_before,
+        "layer_counts_after": layer_counts_after,
+        "stage2_role_counts": stage2["stage2_role"].value_counts().to_dict(),
+        "target_layer_counts": stage2["target_layer"].value_counts().to_dict(),
+    }
+    (out / "manifest.json").write_text(json.dumps(manifest, indent=2, ensure_ascii=False))
+    print(json.dumps(manifest, indent=2, ensure_ascii=False))
+
+
+if __name__ == "__main__":
+    main()