Update run_analysis task4 compatibility

analysis/run_analysis.py

@@ -275,43 +275,71 @@ def run_task4(phase, model, src_tok, tgt_tok, device, cfg,
     print(f"  TASK 4 — Step Ablation  (phase={phase})")
     print("="*65)
 
-    from analysis.step_ablation import (generate_ablation_configs,
-        run_ablation_analysis, plot_ablation_3d, run_adversarial_test)
+    import analysis.step_ablation as step_ablation
+
+    # Legacy API
+    has_legacy = all(hasattr(step_ablation, fn) for fn in [
+        "generate_ablation_configs", "run_ablation_analysis", "plot_ablation_3d"
+    ])
+
+    # New API
+    has_new = hasattr(step_ablation, "run_task4")
 
     if phase == "generate_configs":
-        print("  Generating ablation configs...")
-        generate_ablation_configs(output_dir="ablation_configs")
-        print("\n  NEXT STEPS:")
-        print("  1. bash ablation_configs/train_all.sh")
-        print("  2. python analysis/run_analysis.py --task 4 --phase analyze")
-
-    elif phase == "analyze":
-        # Check which models exist
+        if has_legacy:
+            print("  Generating ablation configs...")
+            step_ablation.generate_ablation_configs(output_dir="ablation_configs")
+            print("\n  NEXT STEPS:")
+            print("  1. bash ablation_configs/train_all.sh")
+            print("  2. python analysis/run_analysis.py --task 4 --phase analyze")
+            return
+        print("  This step_ablation version does not expose config generation helpers.")
+        print("  Use your latest ablation training script/config pipeline directly.")
+        return
+
+    if phase == "analyze":
         existing = [T for T in [4, 8, 16, 32, 64]
                     if os.path.exists(f"ablation_results/T{T}/best_model.pt")]
         if not existing:
             print("  No ablation models found at ablation_results/T*/best_model.pt")
-            print("  Run: python analysis/run_analysis.py --task 4 --phase generate_configs")
-            print("  Then: bash ablation_configs/train_all.sh")
             return
-
         print(f"  Found models for T={existing}")
-        results = run_ablation_analysis(
-            ablation_dir="ablation_results", base_cfg=cfg,
-            src_list=src_list[:200], ref_list=ref_list[:200],
-            tgt_tokenizer=tgt_tok, device=device,
-            output_dir=OUTPUT_DIR)
-        plot_ablation_3d(results,
-            save_path=os.path.join(OUTPUT_DIR, "task4_ablation_3d.png"))
-
-    # Adversarial robustness always runs on existing model (no retraining)
-    print("\n  Running adversarial robustness test...")
-    inp_texts = [src_tok.decode([x for x in s[0].tolist() if x > 4])
-                 for s in src_list[:50]]
-    run_adversarial_test(
-        model, src_tok, tgt_tok,
-        test_inputs=inp_texts, test_refs=ref_list[:50],
-        device=device, output_dir=OUTPUT_DIR)
+
+        if has_legacy:
+            results = step_ablation.run_ablation_analysis(
+                ablation_dir="ablation_results", base_cfg=cfg,
+                src_list=src_list[:200], ref_list=ref_list[:200],
+                tgt_tokenizer=tgt_tok, device=device,
+                output_dir=OUTPUT_DIR)
+            step_ablation.plot_ablation_3d(
+                results, save_path=os.path.join(OUTPUT_DIR, "task4_ablation_3d.png"))
+        elif has_new:
+            from inference import load_model as _load_model
+            models = {}
+            for T in existing:
+                ckpt = f"ablation_results/T{T}/best_model.pt"
+                cfg_t = copy.deepcopy(cfg)
+                cfg_t["model"]["diffusion_steps"] = T
+                cfg_t["inference"]["num_steps"] = T
+                m_t, _ = _load_model(ckpt, cfg_t, device)
+                m_t.eval()
+                models[T] = m_t
+            knee_t = step_ablation.run_task4(
+                models, src_list[:200], ref_list[:200], tgt_tok)
+            print(f"  New pipeline suggested optimal T={knee_t}")
+        else:
+            print("  Unsupported step_ablation API; please sync analysis/step_ablation.py")
+            return
+
+    # Optional adversarial robustness (legacy helper only)
+    if hasattr(step_ablation, "run_adversarial_test"):
+        print("\n  Running adversarial robustness test...")
+        inp_texts = [src_tok.decode([x for x in s[0].tolist() if x > 4])
+                     for s in src_list[:50]]
+        step_ablation.run_adversarial_test(
+            model, src_tok, tgt_tok,
+            test_inputs=inp_texts, test_refs=ref_list[:50],
+            device=device, output_dir=OUTPUT_DIR)
 
 
 # ── Task 5 ────────────────────────────────────────────────────────────