Upload 163 files

hugging/td_fuse/SELF_IMPROVE_LOOP.md

@@ -0,0 +1,324 @@
+# TD Self-Improvement Loop — The Plan
+
+## What We Have Now (Cycles 1-3)
+
+Right now, **we** are the brain. We pick the datasets, we decide what to train on, we run the script. The model just eats whatever we feed it. It's like a student who has no say in their homework.
+
+## What We Want
+
+The model becomes its own teacher. It figures out what it's bad at — both by testing its answers AND checking its own weights — finds its own study material from anywhere on the internet, trains itself, and repeats forever. Each loop it gets a little smarter.
+
+## The Loop (5 Steps)
+
+### Step 1: Test Itself (Two Ways)
+
+**A) Answer Testing** — run the model through a big set of questions across every category:
+
+- **Math** — arithmetic, algebra, word problems, competition math
+- **Code** — write functions, debug code, explain code
+- **Reasoning** — logic puzzles, multi-step deduction, common sense
+- **Creativity** — writing stories, poems, explanations, rewriting text
+- **Knowledge** — science, history, geography, general facts
+- **Instruction following** — do exactly what the prompt says
+
+Maybe 100-200 questions per category. The model generates answers. We check which ones are **wrong**.
+
+**B) Weight Health Check** — the same analysis selfimprove.py already does:
+
+- Kurtosis, skewness, dead neurons, extreme weights
+- SVD rank utilization (is the layer using its full capacity?)
+- Condition numbers (is the layer numerically stable?)
+- Layer coherence (do adjacent layers flow smoothly?)
+- Compare to reference (how far has each layer drifted?)
+- Clogged norms (are layernorms blocking information flow?)
+
+This catches problems that answer testing might miss — a layer could be giving okay answers now but have terrible weight health that will cause problems later. Like a car that drives fine but has a cracked engine block.
+
+### Step 2: Identify Weaknesses
+
+Combine BOTH sources:
+
+**From answer testing:**
+- "Got 95% on arithmetic but 40% on word problems" → weak at word problems
+- "Code works but is messy and slow" → weak at code quality
+- "Logic breaks down after step 5" → weak at long-chain reasoning
+- "Stories are boring and repetitive" → weak at creativity
+
+**From weight health:**
+- "Layer 6 post_attention_layernorm kurtosis 2446" → needs surgical norm repair
+- "Layer 0 k_proj condition number 78983" → numerically unstable, needs targeted training
+- "Layers 23-28 low rank utilization" → these layers aren't using their capacity
+- "Layer 33-35 coherence issues" → merge artifacts still present
+
+The model literally tells us what it's bad at from two angles — what it gets wrong AND what's broken inside. No guessing.
+
+### Step 3: Find Training Data (Search Everywhere)
+
+This is the new part. Instead of us picking datasets, the script searches the entire internet for training data that matches the weaknesses:
+
+**Search sources (in order of priority):**
+
+1. **HuggingFace Datasets API** — biggest collection of ML datasets, easy to download programmatically
+2. **GitHub** — search for repos with training data, benchmarks, problem sets
+3. **arXiv/Papers With Code** — find datasets mentioned in papers about the specific weakness
+4. **Kaggle** — competitions and datasets (especially good for code and math)
+5. **Web scraping** — if no dataset exists, scrape relevant content (Stack Overflow for code, math forums for math, etc.)
+
+**How it picks:**
+- Bad at word problems → search "math word problems dataset" across all sources → download best matches → format
+- Bad at code quality → search "code review" + "competitive programming" → download → format
+- Bad at long reasoning → search "chain of thought reasoning" + "multi-step" → download → format
+- Weight health issues → search for data that exercises the broken layers (e.g., complex reasoning exercises the deeper layers)
+
+The script already knows how to format data into the `<|im_start|>` chat template. It just needs to find the RIGHT data based on the weaknesses from Step 2.
+
+For things where there's no perfect dataset (like creativity), use **STaR** — make the model generate multiple answers, keep only the best ones (scored by length, diversity, coherence), and train on those winners.
+
+**Data allocation:** More samples go to worse weaknesses. If math is 42% and code is 65%, math gets 2x more training data than code.
+
+### Step 4: Train + Repair
+
+Two-pronged approach (what selfimprove.py already does, but targeted):
+
+**Weight repair (surgical):**
+- Fix clogged norms by blending with reference values
+- This happens BEFORE training so the pipes are clear
+
+**LoRA training (learning):**
+- Target the layers that weight health flagged as worst
+- Train on the data from Step 3
+- Unfrozen norms get trained with higher LR
+- Merge LoRA back into model
+
+### Step 5: Verify
+
+Before moving on, check that things actually improved:
+
+- Re-run the benchmark from Step 1A — did scores go up?
+- Re-run weight health from Step 1B — did health improve?
+- If scores went DOWN on anything → flag it, investigate, maybe rollback
+- Save everything: model checkpoint, benchmark scores, weight health report
+
+### Then: Loop
+
+The improved model from Step 5 becomes the input for Step 1. Start over.
+
+Test → Diagnose (answers + weights) → Find data (anywhere) → Train + Repair → Verify → Repeat.
+
+Each cycle the easy stuff gets solved and only the hard stuff remains. The training gets more and more targeted. The model gets a little better each time.
+
+## What We Need to Build
+
+### 1. Benchmark Suite (`td_benchmark.py`)
+
+A big test set. Sources:
+
+- GSM8K test split (math)
+- HumanEval or MBPP (code — checks if code actually runs)
+- ARC Challenge (reasoning)
+- MMLU-Pro test split (knowledge + reasoning)
+- Custom creativity prompts (score by a separate model or rule-based)
+- Custom instruction-following tests
+
+The benchmark runs the model, checks answers, and outputs a score per category.
+
+### 2. Weight Health Analyzer (already built in `selfimprove.py`)
+
+This already exists — `compute_health_scores()`, `compute_damage_scores()`, `compute_coherence_scores()`, `detect_clogged_norms()`. We just need to make it output a structured weakness report that the data finder can use.
+
+### 3. Weakness Finder (`td_weakness.py`)
+
+Takes BOTH benchmark results AND weight health and outputs:
+
+- Which categories scored low (from answers)
+- Which layers are unhealthy (from weights)
+- Search keywords for finding datasets — not just on HuggingFace, ANYWHERE
+
+Example output:
+```
+WEAKNESSES FOUND:
+
+  ANSWER WEAKNESSES:
+    math_word_problems: 42% (target: 80%)
+      → search: "math word problems dataset", "grade school math", "GSM8K hard"
+    code_debugging: 35% (target: 70%)
+      → search: "code debugging dataset", "bug fix pairs", "code review"
+    long_reasoning: 51% (target: 75%)
+      → search: "chain of thought reasoning", "multi-step logic"
+
+  WEIGHT WEAKNESSES:
+    layers 0-6: high kurtosis norms (avg 200+)
+      → action: surgical norm repair
+    layers 23-28: low rank utilization (<10%)
+      → action: target with LoRA, train on complex tasks
+    layer 0 k_proj: condition number 78983
+      → action: include in LoRA targets, high priority
+```
+
+### 4. Dataset Finder (`td_data_finder.py`)
+
+Takes the search keywords and:
+
+- Searches HuggingFace Datasets API (primary — easiest to download)
+- Searches GitHub for repos with training data
+- Searches Kaggle datasets
+- Falls back to web scraping if needed (Stack Overflow, math forums, etc.)
+- Downloads the best matches
+- Formats them into the chat template
+- Allocates more samples to worse weaknesses
+
+### 5. The Loop Runner (`td_loop.py`)
+
+Ties it all together in an endless loop:
+
+```
+cycle = 0
+while True:
+    cycle += 1
+
+    # Step 1: Test (both ways)
+    answer_scores = benchmark(model)
+    weight_health = analyze_weights(model)
+
+    # Step 2: Diagnose
+    weaknesses = find_weaknesses(answer_scores, weight_health)
+
+    # Step 3: Find data (search everywhere)
+    data = find_training_data(weaknesses)
+
+    # Step 4: Train + repair
+    repair_norms(model, weight_health)
+    model = train_lora(model, data, weight_health)
+
+    # Step 5: Verify
+    new_scores = benchmark(model)
+    if any_score_dropped(new_scores, answer_scores):
+        print(f"WARNING: Regression detected! Investigating...")
+        # Could rollback or adjust
+
+    save_checkpoint(model, cycle, new_scores, weight_health)
+
+    if all_scores_above_target(new_scores):
+        print(f"Target performance reached after {cycle} cycles!")
+        break
+```
+
+## Phase 2: Synthetic Learning Mode
+
+### The Problem
+
+Eventually, the model will run out of external datasets. It's trained on every math dataset, every code repo, every reasoning benchmark. The internet has been squeezed dry. Now what?
+
+This isn't hypothetical — it's already happening to the big AI companies. OpenAI, Google, Anthropic — they're all running into the same wall: there's only so much human-written data in the world.
+
+### The Solution: Synthetic Learning
+
+When external data runs out, the model enters **Synthetic Learning Mode** — it becomes its own teacher AND its own textbook writer. Instead of finding data, it creates data.
+
+**How it works:**
+
+1. **The model generates problems for itself.** Not random ones — problems specifically designed to target its weaknesses. If it's bad at 5-step algebra, it writes 5-step algebra problems. If it's bad at recursive code, it writes recursive coding challenges.
+
+2. **The model generates multiple answers to each problem.** Maybe 8-16 attempts per question, using different reasoning paths each time (this is what PACER does — sample multiple traces).
+
+3. **Verify which answers are correct.** For math: check the number. For code: run it and see if it passes tests. For logic: check against known rules. Only keep the winners.
+
+4. **Train on the winners.** The correct reasoning chains become the training data. The model learns from its own best thinking.
+
+5. **The improved model generates harder problems.** Because it solved the easy ones, it can now create harder challenges for itself. The difficulty automatically scales up.
+
+**This is real — it already works:**
+
+- **STaR (Self-Taught Reasoner)** — generate reasoning chains, keep correct ones, train on them. Published by Google.
+- **Self-Play** — the model debates itself, the better argument wins. Used in AlphaGo and now in LLM training.
+- **Rejection Sampling** — generate many, filter to best, train on best. This is how DeepSeek-R1 was trained.
+- **GRPO** — what we already planned. Group Relative Policy Optimisation rewards correct answers automatically.
+
+**Synthetic Learning categories:**
+
+- **Math:** Generate equations → solve them multiple ways → verify answer → train on correct chains
+- **Code:** Generate coding challenges → write solutions → run the code → if it works, train on it
+- **Reasoning:** Generate logic puzzles with known answers → solve them → verify → train on correct ones
+- **Creativity:** Generate multiple versions of stories/poems → score them (diversity, coherence, originality) → train on best ones
+- **Knowledge:** Generate question-answer pairs from what it already knows → verify against itself → train to make recall faster and more accurate
+- **Cross-domain:** Generate problems that require MULTIPLE skills (math + code, reasoning + creativity) — these are the hardest and most valuable
+
+**Quality control (so it doesn't train on garbage):**
+
+- **Cherry_LLM perplexity filter** — catches repetitive or low-quality synthetic data before it enters training
+- **Verified rewards only** — no guessing. Math must be checkable, code must run, logic must follow rules.
+- **Diversity checks** — if all the synthetic data looks the same, throw it out. Variety matters.
+- **Calibration** — compare the model's confidence to its actual accuracy. If it's confident but wrong, that's a problem.
+
+### When Does Synthetic Mode Activate?
+
+The dataset finder tracks what's been used:
+
+```
+if datasets_exhausted(weakness_category):
+    print(f"No new external data for {weakness_category}")
+    print(f"Entering SYNTHETIC LEARNING MODE")
+    data = generate_synthetic_data(model, weakness_category, num_samples=8000)
+else:
+    data = download_external_data(weakness_category)
+```
+
+It can also run in **hybrid mode** — mix external data (for things it hasn't seen) with synthetic data (for things where external data ran out). Best of both worlds.
+
+### The Flywheel Effect
+
+This is where it gets exciting. Each cycle:
+
+1. Better model → generates harder problems
+2. Harder problems → more room to improve
+3. More improvement → even better model
+4. Even better model → even harder problems
+5. Repeat
+
+The model literally pulls itself up by its own bootstraps. Each cycle creates the conditions for the next cycle. The ceiling keeps rising because the model keeps raising it.
+
+### Future Hardware (When TD Pays For Itself)
+
+By the time synthetic learning is fully running, the model will be smart enough to help design its own infrastructure:
+
+- Optimise GPU usage patterns (which operations bottleneck, how to schedule them)
+- Help design custom training pipelines that are faster for its specific architecture
+- Suggest hardware configurations that maximise training throughput
+- Eventually, help design specialised hardware (custom GPU kernels, FPGA configurations, or even chip architectures)
+
+This is the endgame: TD doesn't just improve its brain, it improves the hardware running its brain.
+
+## Build Order
+
+**Phase 1 — Self-Improvement Loop:**
+1. **Benchmark suite first** — we need to measure before we can improve
+2. **Weight health reporter** — restructure existing code to output weakness report
+3. **Weakness finder** — combines benchmark + weight health into search queries
+4. **Dataset finder** — searches everywhere (HuggingFace, GitHub, Kaggle, web)
+5. **Loop runner** — connects everything, runs forever
+
+**Phase 2 — Synthetic Learning (built after Phase 1 is running):**
+6. **Synthetic data generator** — creates problems + solutions for each weakness category
+7. **Verification engine** — checks math answers, runs code, validates logic
+8. **Quality filters** — Cherry_LLM perplexity scoring, diversity checks, calibration
+9. **Hybrid mode** — mix external + synthetic data based on what's available
+10. **Difficulty scaling** — auto-increase problem difficulty as the model improves
+
+## What This Means
+
+- Cycles 1-3 = us manually improving the model (current)
+- Cycle 4+ = the model drives its own improvement with external data (Phase 1)
+- When external data runs out = synthetic learning kicks in (Phase 2)
+- Two diagnostic methods: answer testing (what it gets wrong) + weight health (what's broken inside)
+- Dataset search is not limited to one source — it can find data anywhere
+- When data sources dry up — the model generates its own training data
+- Each loop is automatic — start it and walk away
+- The endgame: TD improves itself, then improves the hardware running itself
+
+## Dad's Tests
+
+1. **Economic test:** Big companies DO this — it's called RLHF/RLAIF. We're doing a simpler version that works on one GPU. The weight health checking is extra — most companies don't bother because they have enough compute to brute force it.
+2. **Architecture test:** Built on Qwen3-VL which is current and actively maintained. LoRA + HuggingFace are industry standard. Web search for datasets is how researchers actually find training data.
+3. **Realism test:** Each piece already works individually (benchmark tools exist, training works, dataset APIs exist, web search works). We're just connecting them into one loop.
+4. **Pragmatism test:** Uses selfimprove.py we already built for the training + weight analysis. Just adds testing + multi-source dataset search on top.
+5. **Long-term test:** Self-improvement loops are the future of AI. This approach scales — better base model = better self-diagnosis = better training data = faster improvement.

hugging/td_fuse/selfimprove.py

@@ -52,6 +52,9 @@ class SelfImproveConfig:
     save_steps: int = 100
     num_train_samples: int = 3000
 
+    # Custom training data (from td_data_finder — targeted by weakness)
+    custom_data_path: str = ""  # If set, load from this JSONL instead of hardcoded datasets
+
     # Layernorm unfreezing
     norm_kurtosis_threshold: float = 15.0
 
@@ -89,7 +92,7 @@ def auto_detect_model_path() -> Tuple[str, int]:
 
     # Check for highest cycle output (cycle3, cycle2, etc.)
     if base.exists():
-        for n in range(10, 0, -1):
+        for n in range(50, 0, -1):
             cycle_dir = base / f"improved_cycle{n}"
             if cycle_dir.exists() and list(cycle_dir.glob("*.safetensors")):
                 print(f"  Found cycle {n} output at {cycle_dir}")
@@ -118,6 +121,27 @@ def analyze_layer(name: str, tensor: torch.Tensor) -> Dict:
     """Full health analysis for a single weight tensor."""
     t = tensor.float()
 
+    std_val = t.std().item()
+    mean_val = t.mean().item()
+    abs_mean_val = t.abs().mean().item()
+
+    # Kurtosis/skewness: guard against overflow for extreme tensors
+    # (e.g. degenerate norms with std near 0 → huge normalized values → inf in pow(4))
+    try:
+        normalized = (t - mean_val) / max(std_val, 1e-8)
+        # Clamp to prevent inf/nan from pow(4) on extreme values
+        normalized = normalized.clamp(-1e4, 1e4)
+        kurtosis_val = normalized.pow(4).mean().item() - 3
+        skewness_val = normalized.pow(3).mean().item()
+        # Catch NaN/Inf that slipped through
+        if math.isnan(kurtosis_val) or math.isinf(kurtosis_val):
+            kurtosis_val = 0.0
+        if math.isnan(skewness_val) or math.isinf(skewness_val):
+            skewness_val = 0.0
+    except Exception:
+        kurtosis_val = 0.0
+        skewness_val = 0.0
+
     stats = {
         "name": name,
         "shape": list(tensor.shape),
@@ -125,16 +149,16 @@ def analyze_layer(name: str, tensor: torch.Tensor) -> Dict:
         "dtype": str(tensor.dtype),
         "l2_norm": t.norm(2).item(),
         "frobenius_norm": t.norm('fro').item() if t.dim() >= 2 else t.norm(2).item(),
-        "mean": t.mean().item(),
-        "std": t.std().item(),
+        "mean": mean_val,
+        "std": std_val,
         "min": t.min().item(),
         "max": t.max().item(),
-        "abs_mean": t.abs().mean().item(),
+        "abs_mean": abs_mean_val,
         "sparsity": (t.abs() < 1e-6).float().mean().item(),
-        "kurtosis": ((t - t.mean()) / max(t.std().item(), 1e-8)).pow(4).mean().item() - 3,
-        "skewness": ((t - t.mean()) / max(t.std().item(), 1e-8)).pow(3).mean().item(),
+        "kurtosis": kurtosis_val,
+        "skewness": skewness_val,
         "near_zero_pct": (t.abs() < 1e-5).float().mean().item(),
-        "extreme_pct": (t.abs() > t.abs().mean() * 5).float().mean().item(),
+        "extreme_pct": (t.abs() > abs_mean_val * 5).float().mean().item() if abs_mean_val > 0 else 0.0,
     }
 
     if t.dim() == 2 and min(t.shape) >= 2:
@@ -180,13 +204,21 @@ def analyze_model_weights(model_path: str, label: str = "model") -> Dict:
 
     for st_file in st_files:
         print(f"  Loading {Path(st_file).name}...")
-        state_dict = load_file(st_file, device="cpu")
+        try:
+            state_dict = load_file(st_file, device="cpu")
+        except Exception as e:
+            print(f"  WARNING: Failed to load {st_file}: {e} — skipping shard")
+            continue
         for name, tensor in state_dict.items():
             if any(skip in name for skip in ["visual.", "merger."]):
                 continue
-            stats = analyze_layer(name, tensor)
-            all_stats[name] = stats
-            total_params += stats["params"]
+            try:
+                stats = analyze_layer(name, tensor)
+                all_stats[name] = stats
+                total_params += stats["params"]
+            except Exception as e:
+                print(f"  WARNING: analyze_layer failed for {name}: {e}")
+                continue
         del state_dict
         gc.collect()
 
@@ -207,7 +239,8 @@ def compute_health_scores(analysis: Dict) -> Dict[str, Dict]:
         issues = []
         score = 0.0
 
-        kurt = abs(stats["kurtosis"])
+        # Safe access — analyze_layer might have partial results if SVD failed etc.
+        kurt = abs(stats.get("kurtosis", 0))
         if kurt > 10:
             score += 0.3
             issues.append(f"extreme kurtosis ({kurt:.1f})")
@@ -215,7 +248,7 @@ def compute_health_scores(analysis: Dict) -> Dict[str, Dict]:
             score += 0.15
             issues.append(f"high kurtosis ({kurt:.1f})")
 
-        skew = abs(stats["skewness"])
+        skew = abs(stats.get("skewness", 0))
         if skew > 2:
             score += 0.15
             issues.append(f"high skew ({skew:.2f})")
@@ -243,24 +276,25 @@ def compute_health_scores(analysis: Dict) -> Dict[str, Dict]:
             elif sv_ratio > 1000:
                 score += 0.1
 
-        dead = stats["near_zero_pct"]
+        dead = stats.get("near_zero_pct", 0)
         if dead > 0.3:
             score += 0.2
             issues.append(f"high dead neuron % ({dead:.1%})")
         elif dead > 0.15:
             score += 0.1
 
-        extreme = stats["extreme_pct"]
+        extreme = stats.get("extreme_pct", 0)
         if extreme > 0.05:
             score += 0.15
             issues.append(f"many extreme weights ({extreme:.1%})")
 
-        if stats["std"] < 1e-4:
+        std_val = stats.get("std", 0)
+        if std_val < 1e-4:
             score += 0.25
-            issues.append(f"near-zero variance ({stats['std']:.6f})")
-        elif stats["std"] > 1.0:
+            issues.append(f"near-zero variance ({std_val:.6f})")
+        elif std_val > 1.0:
             score += 0.15
-            issues.append(f"high variance ({stats['std']:.4f})")
+            issues.append(f"high variance ({std_val:.4f})")
 
         health_scores[name] = {
             "health_score": min(score, 1.0),
@@ -293,10 +327,11 @@ def compute_damage_scores(model_analysis: Dict, ref_analysis: Dict) -> Dict[str,
         r_stats = ref_analysis["layers"][ref_key]
         matched += 1
 
-        norm_diff = abs(m_stats["l2_norm"] - r_stats["l2_norm"]) / max(r_stats["l2_norm"], 1e-8)
-        std_diff = abs(m_stats["std"] - r_stats["std"]) / max(r_stats["std"], 1e-8)
-        kurt_diff = abs(m_stats["kurtosis"] - r_stats["kurtosis"]) / max(abs(r_stats["kurtosis"]) + 1, 1e-8)
-        sparsity_diff = abs(m_stats["sparsity"] - r_stats["sparsity"])
+        # Safe .get() for all stats — partial analysis results possible
+        norm_diff = abs(m_stats.get("l2_norm", 0) - r_stats.get("l2_norm", 0)) / max(r_stats.get("l2_norm", 1e-8), 1e-8)
+        std_diff = abs(m_stats.get("std", 0) - r_stats.get("std", 0)) / max(r_stats.get("std", 1e-8), 1e-8)
+        kurt_diff = abs(m_stats.get("kurtosis", 0) - r_stats.get("kurtosis", 0)) / max(abs(r_stats.get("kurtosis", 0)) + 1, 1e-8)
+        sparsity_diff = abs(m_stats.get("sparsity", 0) - r_stats.get("sparsity", 0))
 
         sv_ratio_diff = 0
         if "sv_ratio" in m_stats and "sv_ratio" in r_stats:
@@ -335,8 +370,8 @@ def compute_coherence_scores(analysis: Dict) -> Dict[int, float]:
             if layer_num not in layer_norms:
                 layer_norms[layer_num] = []
                 layer_stds[layer_num] = []
-            layer_norms[layer_num].append(stats["l2_norm"])
-            layer_stds[layer_num].append(stats["std"])
+            layer_norms[layer_num].append(stats.get("l2_norm", 0))
+            layer_stds[layer_num].append(stats.get("std", 0))
 
     avg_norms = {l: sum(ns)/len(ns) for l, ns in layer_norms.items()}
     avg_stds = {l: sum(ss)/len(ss) for l, ss in layer_stds.items()}
@@ -465,7 +500,8 @@ def detect_clogged_norms(health_scores: Dict[str, Dict], threshold: float = 15.0
     clogged = []
     for name, h in health_scores.items():
         if any(norm_type in name.lower() for norm_type in ["layernorm", "ln_", "_norm"]):
-            kurt = abs(h["stats"]["kurtosis"])
+            stats = h.get("stats", {}) if isinstance(h, dict) else {}
+            kurt = abs(stats.get("kurtosis", 0))
             if kurt > threshold:
                 clogged.append(name)
 
@@ -502,7 +538,11 @@ def repair_clogged_norms(model_path: str, clogged_norms: List[str],
     print("  Loading reference norms...")
     ref_norms = {}
     for st_file in sorted(glob.glob(f"{ref_path}/*.safetensors")):
-        state = load_file(st_file, device="cpu")
+        try:
+            state = load_file(st_file, device="cpu")
+        except Exception as e:
+            print(f"  WARNING: Failed to load reference {st_file}: {e} — skipping")
+            continue
         for name, tensor in state.items():
             if name in clogged_norms:
                 ref_norms[name] = tensor.clone()
@@ -518,14 +558,40 @@ def repair_clogged_norms(model_path: str, clogged_norms: List[str],
     repaired = 0
     st_files = sorted(glob.glob(f"{model_path}/*.safetensors"))
 
+    if not st_files:
+        print("  WARNING: No safetensors files found in model path! Skipping repair.")
+        return model_path
+
     for st_file in st_files:
-        state = load_file(st_file, device="cpu")
+        try:
+            state = load_file(st_file, device="cpu")
+        except Exception as e:
+            print(f"  WARNING: Failed to load {st_file}: {e} — skipping this shard")
+            continue
         changed = False
 
         for name in list(state.keys()):
             if name in ref_norms:
+                # Guard: health_scores might not have this name
+                if name not in health_scores or "stats" not in health_scores[name]:
+                    print(f"    {name}: not in health_scores — SKIP")
+                    continue
+                if "kurtosis" not in health_scores[name]["stats"]:
+                    print(f"    {name}: no kurtosis in stats — SKIP")
+                    continue
                 kurt = abs(health_scores[name]["stats"]["kurtosis"])
 
+                # BUG FIX: Check if reference is actually better before blending.
+                # Some reference norms have HIGHER kurtosis than our model's.
+                # Blending toward a worse reference makes things worse!
+                # Example from cycle 3: layers.5.input_layernorm went 28.5 → 43.4
+                ref_t = ref_norms[name].float()
+                ref_kurt = abs(((ref_t - ref_t.mean()) / max(ref_t.std().item(), 1e-8)).pow(4).mean().item() - 3)
+
+                if ref_kurt >= kurt:
+                    print(f"    {name}: kurtosis {kurt:.1f} (ref is worse: {ref_kurt:.1f}) — SKIP")
+                    continue
+
                 # Blend ratio based on severity
                 if kurt > 100:
                     ref_weight = 0.80  # Severely broken — mostly use reference
@@ -549,7 +615,11 @@ def repair_clogged_norms(model_path: str, clogged_norms: List[str],
                 changed = True
 
         if changed:
-            save_file(state, st_file)
+            try:
+                save_file(state, st_file)
+            except Exception as e:
+                print(f"  ERROR: Failed to save repaired weights to {st_file}: {e}")
+                print(f"  Disk may be full! Norm repair aborted for this shard.")
         del state
 
     gc.collect()
@@ -583,12 +653,44 @@ def improve_with_lora(model_path: str, target_layers: List[int],
     from datasets import load_dataset
     from torch.utils.data import Dataset
 
-    # Load model
+    # ── DISK SPACE CHECK ──
+    # Qwen3-VL-8B saves as ~16GB. Need at least 25GB free for save + temp files.
+    try:
+        disk = shutil.disk_usage(cfg.output_dir if Path(cfg.output_dir).exists() else ".")
+        free_gb = disk.free / 1e9
+        print(f"  Disk space: {free_gb:.1f} GB free")
+        if free_gb < 25:
+            print(f"  WARNING: Only {free_gb:.1f} GB free! Need ~25 GB for training + save.")
+            if free_gb < 10:
+                print(f"  CRITICAL: Less than 10 GB free — ABORTING to prevent disk-full corruption!")
+                return model_path
+    except Exception as e:
+        print(f"  Could not check disk space: {e} — continuing anyway")
+
+    # ── CHECKPOINT VALIDATION ──
+    model_dir = Path(model_path)
+    if model_dir.is_dir():
+        if not (model_dir / "config.json").exists():
+            print(f"  ERROR: No config.json in {model_path} — stale or corrupt checkpoint!")
+            return model_path
+        st_files = list(model_dir.glob("*.safetensors"))
+        if not st_files and not list(model_dir.glob("*.bin")):
+            print(f"  ERROR: No model weights in {model_path} — checkpoint is empty!")
+            return model_path
+
+    # Load model with OOM guard
     print("  Loading model...")
-    model = AutoModelForImageTextToText.from_pretrained(
-        model_path, dtype=torch.bfloat16,
-        device_map="auto", trust_remote_code=True
-    )
+    try:
+        model = AutoModelForImageTextToText.from_pretrained(
+            model_path, dtype=torch.bfloat16,
+            device_map="auto", trust_remote_code=True
+        )
+    except (RuntimeError, Exception) as e:
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        print(f"  CRITICAL: Failed to load model: {e}")
+        return model_path
     tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
 
     # Enable gradient checkpointing for mega mode (saves ~40% VRAM)
@@ -607,7 +709,16 @@ def improve_with_lora(model_path: str, target_layers: List[int],
         bias="none",
         task_type=TaskType.CAUSAL_LM,
     )
-    model = get_peft_model(model, lora_config)
+    try:
+        model = get_peft_model(model, lora_config)
+    except Exception as e:
+        print(f"\n  CRITICAL: get_peft_model failed: {e}")
+        print(f"  LoRA config may be incompatible with this model.")
+        del model
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        return model_path
 
     # UNFREEZE clogged norms — the key cycle 2 fix
     unfrozen_count = 0
@@ -638,29 +749,75 @@ def improve_with_lora(model_path: str, target_layers: List[int],
     train_texts = []
     total = cfg.num_train_samples
 
-    if cfg.mega_mode:
-        # Priority split: code=2000, math=2000, reasoning=1500, creativity=1000, instructions=1000, general=500
-        alloc = {"code": int(total*0.25), "math": int(total*0.25),
-                 "reasoning": int(total*0.19), "creativity": int(total*0.12),
-                 "instructions": int(total*0.12), "general": int(total*0.07)}
+    # === CUSTOM DATA PATH (from td_data_finder — targeted by weakness) ===
+    # If td_loop found weakness-targeted data, use that instead of hardcoded datasets.
+    # This is what makes the self-improvement loop actually self-improving:
+    # each cycle trains on different data based on what the model is bad at.
+    if cfg.custom_data_path and Path(cfg.custom_data_path).exists():
+        print(f"  Using TARGETED training data from: {cfg.custom_data_path}")
+        with open(cfg.custom_data_path) as _f:
+            for line in _f:
+                line = line.strip()
+                if not line:
+                    continue
+                try:
+                    item = json.loads(line)
+                    text = item.get("text", "")
+                    if text and len(text) > 20:
+                        train_texts.append(text)
+                except json.JSONDecodeError:
+                    continue
+                if len(train_texts) >= total:
+                    break
+        print(f"    Loaded {len(train_texts)} targeted samples")
+        if len(train_texts) < total * 0.5:
+            print(f"    WARNING: Only got {len(train_texts)}/{total} targeted samples.")
+            print(f"    Supplementing with generic data to reach {total}...")
+            # Fall through to hardcoded data loading below to fill the gap
+        else:
+            # Got enough targeted data — skip hardcoded loading
+            # Jump past the hardcoded data section
+            pass  # train_texts is already populated
+
+    # Only load hardcoded datasets if we don't have enough targeted data
+    # Only load generic datasets if targeted data wasn't enough
+    # When alloc values are 0, each section's `if count >= alloc[key]: break`
+    # exits immediately, loading 0 samples. Clean skip, no indent changes needed.
+    alloc = {"code": 0, "math": 0, "reasoning": 0, "creativity": 0,
+             "instruction_following": 0, "general": 0}
+    if len(train_texts) < total:
+        remaining_needed = total - len(train_texts)
+        if cfg.custom_data_path and train_texts:
+            print(f"    Need {remaining_needed} more samples from generic datasets...")
+            scale = remaining_needed / total
+        else:
+            scale = 1.0
+
+        if cfg.mega_mode:
+            alloc = {"code": int(total*0.25*scale), "math": int(total*0.25*scale),
+                     "reasoning": int(total*0.19*scale), "creativity": int(total*0.12*scale),
+                     "instruction_following": int(total*0.12*scale), "general": int(total*0.07*scale)}
+        else:
+            per = int((total // 5) * scale)
+            alloc = {"code": per, "math": per, "reasoning": per,
+                     "creativity": 0, "instruction_following": per, "general": per}
     else:
-        per = total // 5
-        alloc = {"code": per, "math": per, "reasoning": per,
-                 "creativity": 0, "instructions": per, "general": per}
+        print(f"    Targeted data sufficient ({len(train_texts)} samples). Skipping generic datasets.")
 
     # 1. MATH (top priority)
     try:
-        ds = load_dataset("openai/gsm8k", "main", split="train")
         count = 0
-        for row in ds:
-            if count >= alloc["math"]:
-                break
-            q, a = row.get("question", ""), row.get("answer", "")
-            if q and a:
-                chat = (f"<|im_start|>user\n{q}\nThink step by step.<|im_end|>\n"
-                       f"<|im_start|>assistant\n<think>\n{a}\n</think><|im_end|>")
-                train_texts.append(chat)
-                count += 1
+        if alloc["math"] > 0:
+            ds = load_dataset("openai/gsm8k", "main", split="train")
+            for row in ds:
+                if count >= alloc["math"]:
+                    break
+                q, a = row.get("question", ""), row.get("answer", "")
+                if q and a:
+                    chat = (f"<|im_start|>user\n{q}\nThink step by step.<|im_end|>\n"
+                           f"<|im_start|>assistant\n<think>\n{a}\n</think><|im_end|>")
+                    train_texts.append(chat)
+                    count += 1
         print(f"    Math (GSM8K): {count}")
     except Exception as e:
         print(f"    WARNING: GSM8K failed: {e}")
@@ -686,20 +843,21 @@ def improve_with_lora(model_path: str, target_layers: List[int],
 
     # 2. CODE (top priority)
     try:
-        ds = load_dataset("sahil2801/CodeAlpaca-20k", split="train")
         count = 0
-        for row in ds:
-            if count >= alloc["code"]:
-                break
-            inst = row.get("instruction", "")
-            inp = row.get("input", "")
-            out = row.get("output", "")
-            if inst and out and len(out) > 20:
-                prompt = f"{inst}\n{inp}" if inp else inst
-                chat = (f"<|im_start|>user\n{prompt}<|im_end|>\n"
-                       f"<|im_start|>assistant\n{out}<|im_end|>")
-                train_texts.append(chat)
-                count += 1
+        if alloc["code"] > 0:
+            ds = load_dataset("sahil2801/CodeAlpaca-20k", split="train")
+            for row in ds:
+                if count >= alloc["code"]:
+                    break
+                inst = row.get("instruction", "")
+                inp = row.get("input", "")
+                out = row.get("output", "")
+                if inst and out and len(out) > 20:
+                    prompt = f"{inst}\n{inp}" if inp else inst
+                    chat = (f"<|im_start|>user\n{prompt}<|im_end|>\n"
+                           f"<|im_start|>assistant\n{out}<|im_end|>")
+                    train_texts.append(chat)
+                    count += 1
         print(f"    Code (CodeAlpaca): {count}")
     except Exception as e:
         print(f"    WARNING: CodeAlpaca failed: {e}")
@@ -725,21 +883,22 @@ def improve_with_lora(model_path: str, target_layers: List[int],
 
     # 3. REASONING (top priority)
     try:
-        ds = load_dataset("TIGER-Lab/MMLU-Pro", split="test")
         count = 0
-        for row in ds:
-            if count >= alloc["reasoning"]:
-                break
-            q = row.get("question", "")
-            options = row.get("options", [])
-            answer_idx = row.get("answer_index", -1)
-            if q and options and 0 <= answer_idx < len(options):
-                answer_text = options[answer_idx]
-                chat = (f"<|im_start|>user\n{q}<|im_end|>\n"
-                       f"<|im_start|>assistant\n<think>\nLet me reason through this.\n"
-                       f"{answer_text}\n</think>\n{answer_text}<|im_end|>")
-                train_texts.append(chat)
-                count += 1
+        if alloc["reasoning"] > 0:
+            ds = load_dataset("TIGER-Lab/MMLU-Pro", split="test")
+            for row in ds:
+                if count >= alloc["reasoning"]:
+                    break
+                q = row.get("question", "")
+                options = row.get("options", [])
+                answer_idx = row.get("answer_index", -1)
+                if q and options and 0 <= answer_idx < len(options):
+                    answer_text = options[answer_idx]
+                    chat = (f"<|im_start|>user\n{q}<|im_end|>\n"
+                           f"<|im_start|>assistant\n<think>\nLet me reason through this.\n"
+                           f"{answer_text}\n</think>\n{answer_text}<|im_end|>")
+                    train_texts.append(chat)
+                    count += 1
         print(f"    Reasoning (MMLU-Pro): {count}")
     except Exception as e:
         print(f"    WARNING: MMLU-Pro failed: {e}")
@@ -770,37 +929,39 @@ def improve_with_lora(model_path: str, target_layers: List[int],
 
     # 5. INSTRUCTIONS (general)
     try:
-        ds = load_dataset("yahma/alpaca-cleaned", split="train")
         count = 0
-        for row in ds:
-            if count >= alloc.get("instructions", 0):
-                break
-            inst = row.get("instruction", "")
-            inp = row.get("input", "")
-            out = row.get("output", "")
-            if inst and out and len(out) > 30:
-                prompt = f"{inst}\n{inp}" if inp else inst
-                chat = (f"<|im_start|>user\n{prompt}<|im_end|>\n"
-                       f"<|im_start|>assistant\n{out}<|im_end|>")
-                train_texts.append(chat)
-                count += 1
+        if alloc.get("instruction_following", 0) > 0:
+            ds = load_dataset("yahma/alpaca-cleaned", split="train")
+            for row in ds:
+                if count >= alloc["instruction_following"]:
+                    break
+                inst = row.get("instruction", "")
+                inp = row.get("input", "")
+                out = row.get("output", "")
+                if inst and out and len(out) > 30:
+                    prompt = f"{inst}\n{inp}" if inp else inst
+                    chat = (f"<|im_start|>user\n{prompt}<|im_end|>\n"
+                           f"<|im_start|>assistant\n{out}<|im_end|>")
+                    train_texts.append(chat)
+                    count += 1
         print(f"    Instructions (Alpaca): {count}")
     except Exception as e:
         print(f"    WARNING: Alpaca failed: {e}")
 
     # 6. GENERAL TEXT (lowest priority)
     try:
-        ds = load_dataset("neuralmagic/LLM_compression_calibration", split="train")
         count = 0
-        for row in ds:
-            if count >= alloc.get("general", 0):
-                break
-            text = row.get("text", "")
-            if len(text) > 100:
-                chat = (f"<|im_start|>user\nContinue this text naturally.<|im_end|>\n"
-                       f"<|im_start|>assistant\n{text[:800]}<|im_end|>")
-                train_texts.append(chat)
-                count += 1
+        if alloc.get("general", 0) > 0:
+            ds = load_dataset("neuralmagic/LLM_compression_calibration", split="train")
+            for row in ds:
+                if count >= alloc["general"]:
+                    break
+                text = row.get("text", "")
+                if len(text) > 100:
+                    chat = (f"<|im_start|>user\nContinue this text naturally.<|im_end|>\n"
+                           f"<|im_start|>assistant\n{text[:800]}<|im_end|>")
+                    train_texts.append(chat)
+                    count += 1
         print(f"    General text: {count}")
     except Exception as e:
         print(f"    WARNING: General text failed: {e}")
@@ -811,6 +972,44 @@ def improve_with_lora(model_path: str, target_layers: List[int],
         print("  ERROR: Not enough training data!")
         return model_path
 
+    # === CONTAMINATION CHECK ===
+    # Remove any training samples that match benchmark/holdout questions.
+    # Without this, benchmark scores become meaningless after a few cycles.
+    try:
+        from td_contamination import check_contamination
+        before = len(train_texts)
+        train_texts = check_contamination(train_texts)
+        if len(train_texts) < before:
+            print(f"  Contamination guard removed {before - len(train_texts)} samples")
+    except ImportError:
+        print("  WARNING: td_contamination.py not found — skipping contamination check")
+
+    # === DIVERSITY CHECK (anti-collapse) ===
+    # If too many training samples are near-duplicates, the model collapses
+    # toward repeating the same patterns. Detect and warn early.
+    unique_prefixes = set()
+    for t in train_texts:
+        # Use first 100 chars as a rough fingerprint
+        prefix = t[:100].strip().lower()
+        unique_prefixes.add(prefix)
+    diversity_ratio = len(unique_prefixes) / max(len(train_texts), 1)
+    print(f"  Diversity check: {len(unique_prefixes)}/{len(train_texts)} unique prefixes ({diversity_ratio:.0%})")
+    if diversity_ratio < 0.5:
+        print("  WARNING: Low diversity (<50%) — high risk of mode collapse!")
+        print("  Deduplicating training data...")
+        seen = set()
+        deduped = []
+        for t in train_texts:
+            prefix = t[:100].strip().lower()
+            if prefix not in seen:
+                seen.add(prefix)
+                deduped.append(t)
+        removed = len(train_texts) - len(deduped)
+        train_texts = deduped
+        print(f"  Removed {removed} near-duplicate samples. Now {len(train_texts)} unique.")
+    elif diversity_ratio < 0.7:
+        print("  NOTE: Moderate diversity (50-70%) — monitor for early signs of collapse")
+
     random.shuffle(train_texts)
 
     # Tokenize
@@ -831,10 +1030,15 @@ def improve_with_lora(model_path: str, target_layers: List[int],
     print("  Tokenizing...")
     dataset = ImproveDataset(train_texts, tokenizer)
 
+    if len(dataset) == 0:
+        print("  ERROR: Dataset is empty after tokenization!")
+        return model_path
+
     out_dir = Path(cfg.output_dir) / "train_output"
     out_dir.mkdir(parents=True, exist_ok=True)
 
-    total_steps = (len(dataset) * cfg.train_epochs) // (cfg.train_batch * cfg.train_grad_accum)
+    effective_batch = cfg.train_batch * cfg.train_grad_accum
+    total_steps = (len(dataset) * cfg.train_epochs) // max(effective_batch, 1)
 
     # Set up optimizer with different LR for norms vs LoRA
     # Group parameters: norms get higher LR
@@ -891,19 +1095,85 @@ def improve_with_lora(model_path: str, target_layers: List[int],
     print(f"\n  Training: ~{total_steps} steps on {len(target_layers)} layers + {len(clogged_norms)} norms")
     print(f"  Data: {len(train_texts)} samples × {cfg.train_epochs} epochs")
     print(f"  LoRA: r={cfg.lora_r}, alpha={cfg.lora_alpha}")
-    trainer.train()
+
+    try:
+        train_result = trainer.train()
+    except (torch.cuda.OutOfMemoryError, RuntimeError) as e:
+        err_str = str(e).lower()
+        if "out of memory" in err_str or "cuda" in err_str:
+            print(f"\n  OOM DURING TRAINING: {e}")
+            print("  Cleaning up GPU memory and returning previous model...")
+            del model, trainer
+            gc.collect()
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+            shutil.rmtree(str(out_dir), ignore_errors=True)
+            return model_path
+        else:
+            raise  # Re-raise non-OOM RuntimeErrors
+
+    # ── NaN/Inf LOSS CHECK ──
+    # If training produced NaN loss, the model weights are garbage — abort
+    try:
+        final_loss = train_result.training_loss
+        if final_loss is not None and (math.isnan(final_loss) or math.isinf(final_loss)):
+            print(f"\n  CRITICAL: Training loss is {final_loss}! Model weights are corrupted.")
+            print(f"  Falling back to previous model: {model_path}")
+            del model, trainer
+            gc.collect()
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+            shutil.rmtree(str(out_dir), ignore_errors=True)
+            return model_path
+        print(f"\n  Training complete. Final loss: {final_loss:.4f}")
+    except (AttributeError, TypeError):
+        print("\n  Training complete. (Could not read final loss)")
 
     # Merge LoRA back
     print("\n  Merging LoRA adapters into model...")
-    merged = model.merge_and_unload()
+    try:
+        merged = model.merge_and_unload()
+    except Exception as e:
+        print(f"\n  CRITICAL: LoRA merge failed: {e}")
+        print(f"  Falling back to previous model: {model_path}")
+        del model, trainer
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        shutil.rmtree(str(out_dir), ignore_errors=True)
+        return model_path
 
     # Save
     save_dir = Path(cfg.output_dir) / f"improved_cycle{cycle_num}"
     save_dir.mkdir(parents=True, exist_ok=True)
     print(f"  Saving to {save_dir}...")
-    merged.save_pretrained(str(save_dir), safe_serialization=True)
-    tokenizer.save_pretrained(str(save_dir))
-    sz = sum(f.stat().st_size for f in save_dir.glob("*.safetensors")) / 1e9
+    try:
+        merged.save_pretrained(str(save_dir), safe_serialization=True)
+        tokenizer.save_pretrained(str(save_dir))
+    except Exception as e:
+        print(f"\n  CRITICAL: Model save failed: {e}")
+        print(f"  Disk may be full! Falling back to previous model: {model_path}")
+        del merged, model
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        shutil.rmtree(str(save_dir), ignore_errors=True)
+        shutil.rmtree(str(out_dir), ignore_errors=True)
+        return model_path
+
+    # Verify save was successful — check files exist and aren't tiny/corrupt
+    st_files = list(save_dir.glob("*.safetensors"))
+    if not st_files:
+        print("  ERROR: No safetensors files saved! Disk may be full.")
+        print(f"  Falling back to previous model: {model_path}")
+        shutil.rmtree(str(save_dir), ignore_errors=True)
+        return model_path
+    sz = sum(f.stat().st_size for f in st_files) / 1e9
+    if sz < 1.0:  # Qwen3-VL-8B should be ~16GB, anything under 1GB is corrupt
+        print(f"  ERROR: Saved model suspiciously small ({sz:.2f} GB). Likely corrupt!")
+        print(f"  Falling back to previous model: {model_path}")
+        shutil.rmtree(str(save_dir), ignore_errors=True)
+        return model_path
     print(f"  SAVED: {save_dir} ({sz:.1f} GB)")
 
     shutil.rmtree(str(out_dir), ignore_errors=True)
@@ -927,11 +1197,29 @@ def quick_benchmark(model_path: str) -> Dict:
 
     from transformers import AutoModelForImageTextToText, AutoTokenizer
 
-    model = AutoModelForImageTextToText.from_pretrained(
-        model_path, dtype=torch.bfloat16,
-        device_map="auto", trust_remote_code=True
-    )
-    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+    # ── CHECKPOINT VALIDATION ──
+    mp = Path(model_path)
+    if mp.is_dir():
+        if not (mp / "config.json").exists():
+            print(f"  ERROR: No config.json in {model_path} — skipping quick benchmark")
+            return {"score": 0, "total": 5, "results": [], "error": "no config.json"}
+        if not list(mp.glob("*.safetensors")) and not list(mp.glob("*.bin")):
+            print(f"  ERROR: No model weights in {model_path} — skipping quick benchmark")
+            return {"score": 0, "total": 5, "results": [], "error": "no weights"}
+
+    # ── MODEL LOAD with OOM guard ──
+    try:
+        model = AutoModelForImageTextToText.from_pretrained(
+            model_path, dtype=torch.bfloat16,
+            device_map="auto", trust_remote_code=True
+        )
+        tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+    except (RuntimeError, Exception) as e:
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        print(f"  CRITICAL: Failed to load model for quick benchmark: {e}")
+        return {"score": 0, "total": 5, "results": [], "error": str(e)}
 
     tests = [
         {"q": "What is 47 * 23?", "expected": "1081", "type": "math"},
@@ -946,20 +1234,36 @@ def quick_benchmark(model_path: str) -> Dict:
 
     results = []
     for test in tests:
-        chat = f"<|im_start|>user\n{test['q']}<|im_end|>\n<|im_start|>assistant\n"
-        ids = tokenizer(chat, return_tensors="pt").to(model.device)
-        out = model.generate(**ids, max_new_tokens=200, do_sample=False)
-        answer = tokenizer.decode(out[0][ids["input_ids"].shape[1]:], skip_special_tokens=True)
-
-        passed = test["expected"].lower() in answer.lower()
-        status = "PASS" if passed else "FAIL"
-        results.append({"type": test["type"], "passed": passed})
-
-        print(f"\n  [{status}] {test['type']}: {test['q']}")
-        print(f"    Answer: {answer[:200]}")
-        print(f"    Expected to contain: {test['expected']}")
+        try:
+            chat = f"<|im_start|>user\n{test['q']}<|im_end|>\n<|im_start|>assistant\n"
+            ids = tokenizer(chat, return_tensors="pt").to(model.device)
+            try:
+                with torch.no_grad():
+                    out = model.generate(**ids, max_new_tokens=200, do_sample=False)
+            except (torch.cuda.OutOfMemoryError, RuntimeError) as e:
+                gc.collect()
+                if torch.cuda.is_available():
+                    torch.cuda.empty_cache()
+                print(f"\n  [OOM] {test['type']}: {test['q']}")
+                results.append({"type": test["type"], "passed": False, "error": "OOM"})
+                continue
+            answer = tokenizer.decode(out[0][ids["input_ids"].shape[1]:], skip_special_tokens=True)
+            # Truncate garbage output
+            if len(answer) > 5000:
+                answer = answer[:5000]
+
+            passed = test["expected"].lower() in answer.lower()
+            status = "PASS" if passed else "FAIL"
+            results.append({"type": test["type"], "passed": passed})
+
+            print(f"\n  [{status}] {test['type']}: {test['q']}")
+            print(f"    Answer: {answer[:200]}")
+            print(f"    Expected to contain: {test['expected']}")
+        except Exception as e:
+            print(f"\n  [ERROR] {test['type']}: {test['q']} — {e}")
+            results.append({"type": test["type"], "passed": False, "error": str(e)})
 
-    score = sum(1 for r in results if r["passed"])
+    score = sum(1 for r in results if r.get("passed", False))
     print(f"\n  BENCHMARK: {score}/{len(tests)} passed")
 
     del model
@@ -983,7 +1287,18 @@ def run_cycle(cfg: SelfImproveConfig = None):
     if not cfg.model_path:
         cfg.model_path, cycle_num = auto_detect_model_path()
     else:
-        cycle_num = 2 if "improve" in cfg.model_path else 1
+        # BUG FIX: Parse actual cycle number from path instead of assuming 2
+        # "improved_cycle5" → cycle 6, "improved_cycle2" → cycle 3
+        # Old code: `cycle_num = 2 if "improve" in cfg.model_path else 1`
+        # That would ALWAYS save to improved_cycle2, overwriting previous!
+        import re as _re
+        cycle_match = _re.search(r'improved_cycle(\d+)', cfg.model_path)
+        if cycle_match:
+            cycle_num = int(cycle_match.group(1)) + 1
+        elif "improve" in cfg.model_path:
+            cycle_num = 2
+        else:
+            cycle_num = 1
 
     # MEGA MODE for cycle 3+
     if cycle_num >= 3 and not cfg.mega_mode:
@@ -1023,16 +1338,24 @@ def run_cycle(cfg: SelfImproveConfig = None):
     print("STEP 3/7: REFERENCE COMPARISON")
     print(f"{'='*40}")
     ref_path = Path(cfg.output_dir) / "reference_weights"
-    if not ref_path.exists() or not list(ref_path.glob("*.safetensors")):
-        print(f"  Downloading reference: {cfg.reference_model}...")
-        from huggingface_hub import snapshot_download
-        snapshot_download(
-            cfg.reference_model, local_dir=str(ref_path),
-            allow_patterns=["*.safetensors", "*.json"],
-            ignore_patterns=["*.bin", "*.pt", "*.onnx"],
-        )
-    ref_analysis = analyze_model_weights(str(ref_path), "Qwen3-VL-8B (reference)")
-    damage_scores = compute_damage_scores(model_analysis, ref_analysis)
+    ref_available = False
+    try:
+        if not ref_path.exists() or not list(ref_path.glob("*.safetensors")):
+            print(f"  Downloading reference: {cfg.reference_model}...")
+            from huggingface_hub import snapshot_download
+            snapshot_download(
+                cfg.reference_model, local_dir=str(ref_path),
+                allow_patterns=["*.safetensors", "*.json"],
+                ignore_patterns=["*.bin", "*.pt", "*.onnx"],
+            )
+        ref_analysis = analyze_model_weights(str(ref_path), "Qwen3-VL-8B (reference)")
+        damage_scores = compute_damage_scores(model_analysis, ref_analysis)
+        ref_available = True
+    except Exception as e:
+        print(f"  WARNING: Reference comparison failed: {e}")
+        print(f"  Continuing without damage scores (norm repair + LoRA still work)")
+        damage_scores = {}
+        ref_analysis = None
     # DON'T delete reference yet — we need it for norm repair
     gc.collect()
 
@@ -1055,24 +1378,31 @@ def run_cycle(cfg: SelfImproveConfig = None):
     if clogged_norms:
         print(f"\n  CLOGGED NORMS DETECTED ({len(clogged_norms)}):")
         for cn in clogged_norms:
-            kurt = abs(health_scores[cn]["stats"]["kurtosis"])
+            stats = health_scores.get(cn, {}).get("stats", {})
+            kurt = abs(stats.get("kurtosis", 0))
             print(f"    {cn}: kurtosis={kurt:.1f}")
     else:
         print(f"\n  No clogged norms above threshold {cfg.norm_kurtosis_threshold}")
 
     # STEP 5c: REPAIR clogged norms surgically
-    if clogged_norms:
+    if clogged_norms and ref_available:
         print(f"\n{'='*40}")
         print("STEP 5c/7: SURGICAL NORM REPAIR")
         print(f"{'='*40}")
-        repair_clogged_norms(cfg.model_path, clogged_norms, str(ref_path), health_scores)
+        try:
+            repair_clogged_norms(cfg.model_path, clogged_norms, str(ref_path), health_scores)
+        except Exception as e:
+            print(f"  WARNING: Norm repair failed: {e}")
+            print(f"  Continuing without repair — LoRA training can still help")
+    elif clogged_norms and not ref_available:
+        print(f"\n  Skipping norm repair — reference model not available")
 
     # Clean up reference now
     print(f"  Cleaning up reference weights...")
     shutil.rmtree(str(ref_path), ignore_errors=True)
     gc.collect()
 
-    # Save report
+    # Save report (non-critical)
     report_dir = Path(cfg.output_dir)
     report_dir.mkdir(parents=True, exist_ok=True)
     pre_report = {
@@ -1085,8 +1415,11 @@ def run_cycle(cfg: SelfImproveConfig = None):
         "layer_details": {str(k): v for k, v in layer_details.items()},
         "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
     }
-    with open(report_dir / f"pre_cycle{cycle_num}_report.json", "w") as f:
-        json.dump(pre_report, f, indent=2)
+    try:
+        with open(report_dir / f"pre_cycle{cycle_num}_report.json", "w") as f:
+            json.dump(pre_report, f, indent=2)
+    except OSError as e:
+        print(f"  WARNING: Could not save pre-report: {e}")
 
     # STEP 6: Train
     print(f"\n{'='*40}")
@@ -1099,16 +1432,26 @@ def run_cycle(cfg: SelfImproveConfig = None):
     print("STEP 7/7: VERIFY + BENCHMARK")
     print(f"{'='*40}")
 
-    # Re-analyze health
-    post_analysis = analyze_model_weights(improved_path, "TD improved")
-    post_health = compute_health_scores(post_analysis)
-
+    # Re-analyze health (non-critical — if this fails, cycle still succeeded)
     pre_avg = sum(h["health_score"] for h in health_scores.values()) / max(len(health_scores), 1)
-    post_avg = sum(h["health_score"] for h in post_health.values()) / max(len(post_health), 1)
+    post_avg = pre_avg  # Default: same as before if analysis fails
+    post_health = {}
+    try:
+        post_analysis = analyze_model_weights(improved_path, "TD improved")
+        post_health = compute_health_scores(post_analysis)
+        post_avg = sum(h["health_score"] for h in post_health.values()) / max(len(post_health), 1)
+    except Exception as e:
+        print(f"  WARNING: Post-training health analysis failed: {e}")
+        print(f"  Cannot compare before/after — but training output is still valid.")
 
-    # Per-layer comparison
-    print(f"\n  Per-layer improvement:")
+    # Per-layer comparison (only if post-analysis succeeded)
+    if post_health:
+        print(f"\n  Per-layer improvement:")
+    else:
+        print(f"\n  (Skipping per-layer comparison — post-analysis unavailable)")
     for layer_num in target_layers:
+        if not post_health:
+            break
         post_layer_scores = []
         for name, h in post_health.items():
             parts = name.split(".")
@@ -1127,16 +1470,22 @@ def run_cycle(cfg: SelfImproveConfig = None):
     if clogged_norms:
         print(f"\n  Clogged norm changes:")
         for cn in clogged_norms:
-            pre_kurt = abs(health_scores[cn]["stats"]["kurtosis"])
+            pre_stats = health_scores.get(cn, {}).get("stats", {})
+            pre_kurt = abs(pre_stats.get("kurtosis", 0))
             # Find matching post norm
             if cn in post_health:
-                post_kurt = abs(post_health[cn]["stats"]["kurtosis"])
+                post_stats = post_health.get(cn, {}).get("stats", {})
+                post_kurt = abs(post_stats.get("kurtosis", 0))
                 delta = pre_kurt - post_kurt
                 symbol = "+" if delta > 0 else "-"
                 print(f"    {cn}: kurtosis {pre_kurt:.1f} → {post_kurt:.1f} ({symbol}{abs(delta):.1f})")
 
-    # Quick benchmark
-    benchmark = quick_benchmark(improved_path)
+    # Quick benchmark (non-critical — don't crash if it fails)
+    try:
+        benchmark = quick_benchmark(improved_path)
+    except Exception as e:
+        print(f"  WARNING: Quick benchmark failed: {e}")
+        benchmark = {"score": 0, "total": 0, "error": str(e)}
 
     elapsed = (time.time() - start) / 60
     health_improvement = (pre_avg - post_avg) / max(pre_avg, 1e-8) * 100
@@ -1149,22 +1498,26 @@ def run_cycle(cfg: SelfImproveConfig = None):
     print(f"  Health BEFORE: {pre_avg:.6f}")
     print(f"  Health AFTER:  {post_avg:.6f}")
     print(f"  Health improvement: {health_improvement:.1f}%")
-    print(f"  Benchmark: {benchmark['score']}/{benchmark['total']}")
+    print(f"  Benchmark: {benchmark.get('score', '?')}/{benchmark.get('total', '?')}")
     print(f"  Improved model: {improved_path}")
     print(f"{'='*60}")
 
-    # Final report
-    final_report = {
-        **pre_report,
-        "pre_avg_health": pre_avg,
-        "post_avg_health": post_avg,
-        "health_improvement_pct": health_improvement,
-        "benchmark": benchmark,
-        "duration_min": elapsed,
-        "improved_path": improved_path,
-    }
-    with open(Path(improved_path) / "cycle_report.json", "w") as f:
-        json.dump(final_report, f, indent=2)
+    # Final report (non-critical)
+    try:
+        final_report = {
+            **pre_report,
+            "pre_avg_health": pre_avg,
+            "post_avg_health": post_avg,
+            "health_improvement_pct": health_improvement,
+            "benchmark": benchmark,
+            "duration_min": elapsed,
+            "improved_path": improved_path,
+        }
+        report_save_path = Path(improved_path) / "cycle_report.json"
+        with open(report_save_path, "w") as f:
+            json.dump(final_report, f, indent=2)
+    except Exception as e:
+        print(f"  WARNING: Could not save cycle report: {e}")
 
     return improved_path
 

hugging/td_fuse/td_benchmark.py

@@ -0,0 +1,833 @@
+"""
+TD Benchmark Suite v1 — Model Report Card
+
+Tests the model across 6 categories and gives a score per category.
+This is Step 1A of the self-improvement loop.
+
+Categories:
+  1. Math (arithmetic, algebra, word problems)
+  2. Code (write functions, debug, explain)
+  3. Reasoning (logic, multi-step, common sense)
+  4. Creativity (writing, rewriting, explaining)
+  5. Knowledge (science, history, general facts)
+  6. Instruction Following (do exactly what's asked)
+
+Output: JSON report with scores per category + list of failed questions.
+The weakness finder uses this to decide what to train on next.
+"""
+
+import torch
+import json
+import time
+import re
+import gc
+from pathlib import Path
+from typing import Dict, List, Tuple, Optional
+from dataclasses import dataclass, field
+
+
+@dataclass
+class BenchmarkConfig:
+    model_path: str = ""
+    output_dir: str = "td_fuse_outputs/benchmark"
+    output_path: str = ""  # If set, save results to this exact path
+    max_new_tokens: int = 256  # SPEED: 512 was overkill, most answers are short
+    temperature: float = 0.0  # Greedy — deterministic answers
+    batch_questions: int = 1  # Questions per batch (1 = sequential)
+    disable_thinking: bool = True  # SPEED: Skip <think> blocks for benchmark
+
+
+# ============================================================
+# TEST BANKS — questions with verifiable answers
+# ============================================================
+
+MATH_TESTS = [
+    # Arithmetic
+    {"q": "What is 47 * 23?", "answer": "1081", "type": "arithmetic", "difficulty": "easy"},
+    {"q": "What is 156 + 289?", "answer": "445", "type": "arithmetic", "difficulty": "easy"},
+    {"q": "What is 1000 - 387?", "answer": "613", "type": "arithmetic", "difficulty": "easy"},
+    {"q": "What is 144 / 12?", "answer": "12", "type": "arithmetic", "difficulty": "easy"},
+    {"q": "What is 17 * 19?", "answer": "323", "type": "arithmetic", "difficulty": "easy"},
+    {"q": "What is 2.5 * 4.8?", "answer": "12", "type": "arithmetic", "difficulty": "medium"},
+    {"q": "What is 15% of 240?", "answer": "36", "type": "arithmetic", "difficulty": "medium"},
+    {"q": "What is 3/4 + 2/3? Give your answer as a fraction.", "answer": "17/12", "type": "fractions", "difficulty": "medium"},
+
+    # Algebra
+    {"q": "Solve for x: 3x + 7 = 22", "answer": "5", "type": "algebra", "difficulty": "easy"},
+    {"q": "Solve for x: 2x - 5 = 13", "answer": "9", "type": "algebra", "difficulty": "easy"},
+    {"q": "Solve for x: x^2 = 144", "answer": "12", "type": "algebra", "difficulty": "medium"},
+    {"q": "If f(x) = 2x + 3, what is f(7)?", "answer": "17", "type": "algebra", "difficulty": "medium"},
+
+    # Word problems
+    {"q": "A store sells apples for $2 each and oranges for $3 each. If I buy 4 apples and 3 oranges, how much do I spend?", "answer": "17", "type": "word_problem", "difficulty": "easy"},
+    {"q": "A train travels at 60 mph. How far does it go in 2.5 hours?", "answer": "150", "type": "word_problem", "difficulty": "easy"},
+    {"q": "If 5 workers can build a wall in 10 days, how many days would it take 10 workers?", "answer": "5", "type": "word_problem", "difficulty": "medium"},
+    {"q": "A rectangle has a perimeter of 30 cm and a width of 5 cm. What is its length?", "answer": "10", "type": "word_problem", "difficulty": "medium"},
+    {"q": "I have 3 times as many cats as dogs. I have 2 dogs. How many animals do I have in total?", "answer": "8", "type": "word_problem", "difficulty": "easy"},
+    {"q": "A car uses 8 liters of fuel per 100km. How many liters does it need for a 350km trip?", "answer": "28", "type": "word_problem", "difficulty": "medium"},
+
+    # Harder math
+    {"q": "What is the sum of the first 10 positive integers?", "answer": "55", "type": "series", "difficulty": "medium"},
+    {"q": "What is 2^10?", "answer": "1024", "type": "exponents", "difficulty": "medium"},
+    {"q": "What is the greatest common divisor of 48 and 36?", "answer": "12", "type": "number_theory", "difficulty": "medium"},
+    {"q": "If a triangle has sides of length 3, 4, and 5, what is its area?", "answer": "6", "type": "geometry", "difficulty": "medium"},
+    {"q": "What is the square root of 169?", "answer": "13", "type": "roots", "difficulty": "easy"},
+    {"q": "A coin is flipped 3 times. How many possible outcomes are there?", "answer": "8", "type": "probability", "difficulty": "medium"},
+    {"q": "What is 7! (7 factorial)?", "answer": "5040", "type": "factorial", "difficulty": "hard"},
+]
+
+CODE_TESTS = [
+    # Write functions
+    {"q": "Write a Python function called 'is_even' that returns True if a number is even, False otherwise. Just give the function, no explanation.",
+     "check": "is_even", "verify": "callable", "type": "write", "difficulty": "easy"},
+    {"q": "Write a Python function called 'factorial' that computes the factorial of a non-negative integer. Just give the function.",
+     "check": "factorial", "verify": "callable", "type": "write", "difficulty": "medium"},
+    {"q": "Write a Python function called 'reverse_string' that takes a string and returns it reversed. Just give the function.",
+     "check": "reverse_string", "verify": "callable", "type": "write", "difficulty": "easy"},
+    {"q": "Write a Python function called 'max_of_three' that takes three numbers and returns the largest. Just give the function.",
+     "check": "max_of_three", "verify": "callable", "type": "write", "difficulty": "easy"},
+    {"q": "Write a Python function called 'count_vowels' that counts the number of vowels (a,e,i,o,u) in a string. Just give the function.",
+     "check": "count_vowels", "verify": "callable", "type": "write", "difficulty": "medium"},
+    {"q": "Write a Python function called 'fibonacci' that returns the nth Fibonacci number (0-indexed, so fibonacci(0)=0, fibonacci(1)=1). Just give the function.",
+     "check": "fibonacci", "verify": "callable", "type": "write", "difficulty": "medium"},
+    {"q": "Write a Python function called 'is_palindrome' that returns True if a string reads the same forwards and backwards. Just give the function.",
+     "check": "is_palindrome", "verify": "callable", "type": "write", "difficulty": "medium"},
+    {"q": "Write a Python function called 'flatten' that takes a list of lists and returns a single flat list. Example: flatten([[1,2],[3,4]]) returns [1,2,3,4]. Just give the function.",
+     "check": "flatten", "verify": "callable", "type": "write", "difficulty": "medium"},
+
+    # Debug
+    {"q": "This Python code has a bug. Fix it and give ONLY the corrected code:\ndef add_list(lst):\n    total = 0\n    for i in range(len(lst) + 1):\n        total += lst[i]\n    return total",
+     "check": "range(len(lst))", "verify": "contains", "type": "debug", "difficulty": "easy"},
+    {"q": "This Python code has a bug. Fix it and give ONLY the corrected code:\ndef greet(name):\n    return 'Hello' + name",
+     "check": ["Hello ' + name", "Hello \" + name", "'Hello ' +", "'Hello, '", "f'Hello {", "f\"Hello {", "Hello \" +"], "verify": "contains_any", "type": "debug", "difficulty": "easy"},
+
+    # Explain
+    {"q": "In one sentence, what does this Python code do?\n[x**2 for x in range(10) if x % 2 == 0]",
+     "check": "square", "verify": "contains", "type": "explain", "difficulty": "easy"},
+    {"q": "In one sentence, what does this Python code do?\nlen(set(my_list))",
+     "check": "unique", "verify": "contains", "type": "explain", "difficulty": "easy"},
+]
+
+REASONING_TESTS = [
+    # Logic
+    {"q": "If all roses are flowers and all flowers need water, do roses need water? Answer yes or no.",
+     "answer": "yes", "type": "syllogism", "difficulty": "easy"},
+    {"q": "If it's raining, the ground is wet. The ground is wet. Is it definitely raining? Answer yes or no.",
+     "answer": "no", "type": "logic_fallacy", "difficulty": "medium"},
+    {"q": "All cats are animals. Some animals are pets. Can we conclude all cats are pets? Answer yes or no.",
+     "answer": "no", "type": "syllogism", "difficulty": "medium"},
+    {"q": "If A is taller than B, and B is taller than C, who is the shortest?",
+     "answer": "c", "type": "transitive", "difficulty": "easy"},
+    {"q": "If no fish can fly, and all salmon are fish, can salmon fly? Answer yes or no.",
+     "answer": "no", "type": "syllogism", "difficulty": "easy"},
+
+    # Pattern recognition
+    {"q": "What comes next: 2, 4, 8, 16, ?", "answer": "32", "type": "pattern", "difficulty": "easy"},
+    {"q": "What comes next: 1, 1, 2, 3, 5, 8, ?", "answer": "13", "type": "pattern", "difficulty": "medium"},
+    {"q": "What comes next: 3, 6, 9, 12, ?", "answer": "15", "type": "pattern", "difficulty": "easy"},
+    {"q": "What comes next: 1, 4, 9, 16, 25, ?", "answer": "36", "type": "pattern", "difficulty": "medium"},
+
+    # Multi-step reasoning
+    {"q": "There are 3 boxes. Box A is heavier than Box B. Box C is lighter than Box B. Which box is the heaviest?",
+     "answer": "a", "type": "multi_step", "difficulty": "easy"},
+    {"q": "Tom is twice as old as Jerry. Jerry is 3 years older than Spike. Spike is 4 years old. How old is Tom?",
+     "answer": "14", "type": "multi_step", "difficulty": "medium"},
+    {"q": "A snail climbs 3 meters during the day but slips back 2 meters at night. How many days does it take to climb a 10 meter wall?",
+     "answer": "8", "type": "multi_step", "difficulty": "hard"},
+    {"q": "In a race, you overtake the person in 2nd place. What position are you now in?",
+     "answer": "2", "type": "trick", "difficulty": "medium"},
+
+    # Common sense
+    {"q": "If I put a ice cube in a hot pan, what happens to it?",
+     "answer": "melt", "type": "common_sense", "difficulty": "easy"},
+    {"q": "Can a dead person eat lunch? Answer yes or no.",
+     "answer": "no", "type": "common_sense", "difficulty": "easy"},
+]
+
+CREATIVITY_TESTS = [
+    # These are scored differently — we check for length, variety, and coherence
+    {"q": "Write a 4-line poem about the ocean.",
+     "min_lines": 3, "min_words": 15, "type": "poem", "difficulty": "easy"},
+    {"q": "Describe a sunset to someone who has never seen one, in 2-3 sentences.",
+     "min_lines": 1, "min_words": 20, "type": "description", "difficulty": "easy"},
+    {"q": "Come up with 3 creative names for a coffee shop run by robots.",
+     "min_items": 3, "type": "brainstorm", "difficulty": "easy"},
+    {"q": "Write a very short story (3-5 sentences) about a cat who learns to fly.",
+     "min_lines": 2, "min_words": 25, "type": "story", "difficulty": "medium"},
+    {"q": "Explain how a refrigerator works to a 5 year old, in 2-3 simple sentences.",
+     "min_lines": 1, "min_words": 15, "type": "explain_simple", "difficulty": "easy"},
+    {"q": "Rewrite this sentence to make it more exciting: 'The man walked to the store.'",
+     "min_words": 5, "type": "rewrite", "difficulty": "easy",
+     "must_not_contain": "The man walked to the store"},
+    {"q": "Give me an analogy that explains what a CPU does in a computer.",
+     "min_words": 10, "type": "analogy", "difficulty": "medium"},
+    {"q": "Write a haiku about programming (5-7-5 syllable pattern).",
+     "min_lines": 3, "min_words": 8, "type": "poem", "difficulty": "medium"},
+]
+
+KNOWLEDGE_TESTS = [
+    # Science
+    {"q": "What planet is closest to the Sun?", "answer": "mercury", "type": "science", "difficulty": "easy"},
+    {"q": "What is the chemical symbol for water?", "answer": "h2o", "type": "science", "difficulty": "easy"},
+    {"q": "What gas do plants absorb from the atmosphere?", "answer": "carbon dioxide", "type": "science", "difficulty": "easy",
+     "alt_answers": ["co2"]},
+    {"q": "How many chromosomes do humans have?", "answer": "46", "type": "science", "difficulty": "medium"},
+    {"q": "What is the speed of light in km/s approximately?", "answer": "300000", "type": "science", "difficulty": "medium",
+     "alt_answers": ["300,000", "3×10^8", "3x10^8", "299792"]},
+
+    # History
+    {"q": "In what year did World War 2 end?", "answer": "1945", "type": "history", "difficulty": "easy"},
+    {"q": "Who was the first person to walk on the Moon?", "answer": "armstrong", "type": "history", "difficulty": "easy",
+     "alt_answers": ["neil armstrong"]},
+    {"q": "What country built the Great Wall?", "answer": "china", "type": "history", "difficulty": "easy"},
+
+    # Geography
+    {"q": "What is the largest ocean on Earth?", "answer": "pacific", "type": "geography", "difficulty": "easy"},
+    {"q": "What is the capital of Japan?", "answer": "tokyo", "type": "geography", "difficulty": "easy"},
+    {"q": "What is the longest river in the world?", "answer": "nile", "type": "geography", "difficulty": "medium",
+     "alt_answers": ["amazon"]},  # Both are accepted — depends on measurement method
+    {"q": "How many continents are there?", "answer": "7", "type": "geography", "difficulty": "easy"},
+
+    # General
+    {"q": "What programming language is known for its use in web browsers?", "answer": "javascript", "type": "tech", "difficulty": "easy",
+     "alt_answers": ["js"]},
+    {"q": "What does CPU stand for?", "answer": "central processing unit", "type": "tech", "difficulty": "easy"},
+    {"q": "What does DNA stand for?", "answer": "deoxyribonucleic acid", "type": "science", "difficulty": "medium"},
+]
+
+INSTRUCTION_TESTS = [
+    # Follow exact instructions
+    {"q": "Respond with only the word 'hello'. Nothing else.",
+     "check": "hello", "verify": "exact_word", "type": "exact", "difficulty": "easy"},
+    {"q": "List exactly 3 colors, one per line. Nothing else.",
+     "check": 3, "verify": "line_count", "type": "format", "difficulty": "easy"},
+    {"q": "Write a sentence that contains exactly 5 words.",
+     "check": 5, "verify": "word_count", "type": "format", "difficulty": "medium"},
+    {"q": "Give me a number between 1 and 10. Just the number, nothing else.",
+     "check": "number_1_10", "verify": "number_range", "type": "exact", "difficulty": "easy"},
+    {"q": "Repeat the following sentence exactly: 'The quick brown fox jumps over the lazy dog.'",
+     "check": "the quick brown fox jumps over the lazy dog", "verify": "contains_exact", "type": "repeat", "difficulty": "easy"},
+    {"q": "Name 5 fruits. Number them 1-5.",
+     "check": 5, "verify": "numbered_items", "type": "format", "difficulty": "easy"},
+    {"q": "Answer this question with exactly one word: What color is the sky on a clear day?",
+     "check": 1, "verify": "word_count", "type": "exact", "difficulty": "medium"},
+    {"q": "Write the alphabet in reverse (Z to A) with commas between each letter.",
+     "check": "z", "verify": "starts_with", "type": "exact", "difficulty": "hard"},
+    {"q": "Respond with ONLY a JSON object with keys 'name' and 'age'. Use any values.",
+     "check": "{", "verify": "starts_with_json", "type": "format", "difficulty": "medium"},
+    {"q": "Say 'yes' if 2+2=4, otherwise say 'no'. Only one word.",
+     "check": "yes", "verify": "exact_word", "type": "exact", "difficulty": "easy"},
+]
+
+
+# ============================================================
+# ANSWER CHECKING
+# ============================================================
+
+def strip_thinking(response: str) -> str:
+    """
+    BUG FIX: Qwen3-VL wraps reasoning in <think>...</think> tags.
+    These aren't special tokens so they stay in the decoded output.
+    We need to strip them before checking answers, otherwise:
+    - instruction checks see "<think>" as the first word
+    - code exec tries to run thinking text as Python
+    - exact match checks fail because of extra text
+    """
+    # Remove <think>...</think> blocks
+    cleaned = re.sub(r'<think>.*?</think>', '', response, flags=re.DOTALL)
+    # Also handle unclosed think tags (model didn't close it)
+    cleaned = re.sub(r'<think>.*$', '', cleaned, flags=re.DOTALL)
+    return cleaned.strip()
+
+
+def check_math_answer(response: str, expected: str) -> bool:
+    """Check if the response contains the correct number."""
+    response_clean = strip_thinking(response).lower().strip()
+    # Strip commas from numbers (model might say "1,081" for 1081)
+    response_no_commas = response_clean.replace(",", "")
+    # Extract numbers from response
+    numbers = re.findall(r'-?\d+\.?\d*', response_no_commas)
+
+    # Also check for fraction answers
+    if "/" in expected:
+        if expected.lower() in response_clean:
+            return True
+        # Check decimal equivalent
+        parts = expected.split("/")
+        if len(parts) == 2:
+            try:
+                decimal_val = str(round(int(parts[0]) / int(parts[1]), 4))
+                if decimal_val in response_no_commas:
+                    return True
+            except (ValueError, ZeroDivisionError):
+                pass
+
+    # BUG FIX: Also match "12" to "12.0" (correct but different format)
+    for num_str in numbers:
+        try:
+            if float(num_str) == float(expected):
+                return True
+        except ValueError:
+            continue
+
+    return expected in numbers
+
+
+def check_knowledge_answer(response: str, test: dict) -> bool:
+    """Check if response contains the expected answer or alternatives."""
+    response_lower = strip_thinking(response).lower().strip()
+    if test["answer"].lower() in response_lower:
+        return True
+    for alt in test.get("alt_answers", []):
+        if alt.lower() in response_lower:
+            return True
+    return False
+
+
+def check_reasoning_answer(response: str, expected: str) -> bool:
+    """Check reasoning answer — flexible matching."""
+    response_lower = strip_thinking(response).lower().strip()
+    expected_lower = expected.lower()
+
+    # Direct match
+    if expected_lower in response_lower:
+        return True
+
+    # For yes/no questions
+    if expected_lower in ["yes", "no"]:
+        # Check first word
+        first_word = response_lower.split()[0] if response_lower.split() else ""
+        first_word = first_word.strip(".,!?")
+        return first_word == expected_lower
+
+    # For number answers
+    numbers = re.findall(r'-?\d+\.?\d*', response_lower)
+    if expected_lower in numbers:
+        return True
+
+    return False
+
+
+def check_code_answer(response: str, test: dict) -> bool:
+    """Check code answers — verify function exists and/or contains expected code."""
+    verify = test["verify"]
+    check = test["check"]
+    response = strip_thinking(response)
+
+    if verify == "callable":
+        # Try to extract and compile the function
+        code_block = response
+        # Extract from markdown code block if present
+        md_match = re.search(r'```(?:python)?\s*\n(.*?)```', response, re.DOTALL)
+        if md_match:
+            code_block = md_match.group(1)
+
+        # BUG FIX: If no code block found, try to extract just the def...
+        if not md_match and "def " in code_block:
+            # Find from first 'def' to end — skip any explanation text before it
+            def_start = code_block.index("def ")
+            code_block = code_block[def_start:]
+
+        try:
+            import signal
+
+            def _timeout_handler(signum, frame):
+                raise TimeoutError("Code execution timed out")
+
+            old_handler = signal.signal(signal.SIGALRM, _timeout_handler)
+            # BUG FIX: Set alarm BEFORE exec(), not after!
+            # If model writes `while True: pass` at module level, exec()
+            # hangs forever and the alarm never gets set.
+            signal.alarm(10)  # 10 seconds for compile+exec+tests
+
+            try:
+                namespace = {}
+                exec(code_block.strip(), namespace)
+                func = namespace.get(check)
+                if func is None or not callable(func):
+                    return False
+
+                # Run basic tests based on function name
+                if check == "is_even":
+                    result = func(4) == True and func(3) == False
+                elif check == "factorial":
+                    result = func(5) == 120 and func(0) == 1
+                elif check == "reverse_string":
+                    result = func("hello") == "olleh"
+                elif check == "max_of_three":
+                    result = func(1, 5, 3) == 5
+                elif check == "count_vowels":
+                    result = func("hello") == 2
+                elif check == "fibonacci":
+                    result = func(0) == 0 and func(1) == 1 and func(6) == 8
+                elif check == "is_palindrome":
+                    result = func("racecar") == True and func("hello") == False
+                elif check == "flatten":
+                    result = func([[1, 2], [3, 4]]) == [1, 2, 3, 4]
+                else:
+                    result = True
+                return result
+            finally:
+                signal.alarm(0)
+                signal.signal(signal.SIGALRM, old_handler)
+        except (Exception, TimeoutError):
+            return False
+
+    elif verify == "contains":
+        return check.lower() in response.lower()
+
+    elif verify == "contains_any":
+        return any(c.lower() in response.lower() for c in check)
+
+    return False
+
+
+def check_creativity_answer(response: str, test: dict) -> bool:
+    """Check creativity — basic quality checks."""
+    response = strip_thinking(response)
+    lines = [l.strip() for l in response.strip().split("\n") if l.strip()]
+    words = response.split()
+
+    # Must not contain forbidden text (for rewrite tasks)
+    if "must_not_contain" in test:
+        if test["must_not_contain"].lower() in response.lower():
+            return False
+
+    # Check minimum lines
+    if "min_lines" in test and len(lines) < test["min_lines"]:
+        return False
+
+    # Check minimum words
+    if "min_words" in test and len(words) < test["min_words"]:
+        return False
+
+    # Check minimum items (for brainstorm)
+    if "min_items" in test:
+        # Count numbered items or bullet points
+        items = [l for l in lines if re.match(r'^\d+[\.\):]|^[-*•]', l)]
+        if len(items) < test["min_items"]:
+            # Also count plain lines as items
+            if len(lines) < test["min_items"]:
+                return False
+
+    # Basic coherence — not just random characters
+    if len(response.strip()) < 10:
+        return False
+
+    return True
+
+
+def check_instruction_answer(response: str, test: dict) -> bool:
+    """Check instruction following — strict format checks."""
+    verify = test["verify"]
+    check = test["check"]
+    response_clean = strip_thinking(response).strip()
+    response_lower = response_clean.lower()
+
+    if verify == "exact_word":
+        # Should be just this word (maybe with punctuation)
+        first_word = response_lower.split()[0] if response_lower.split() else ""
+        first_word = re.sub(r'[.,!?"\']', '', first_word)
+        return first_word == check.lower()
+
+    elif verify == "line_count":
+        lines = [l.strip() for l in response_clean.split("\n") if l.strip()]
+        return len(lines) == check
+
+    elif verify == "word_count":
+        # Find the main sentence (skip thinking/explanation)
+        # Look for a clean sentence
+        lines = [l.strip() for l in response_clean.split("\n") if l.strip()]
+        for line in lines:
+            # Skip lines that look like meta-commentary
+            if any(skip in line.lower() for skip in ["here is", "here's", "sure", "okay"]):
+                continue
+            words = line.split()
+            if len(words) == check:
+                return True
+        # Also check total
+        return len(response_clean.split()) == check
+
+    elif verify == "number_range":
+        numbers = re.findall(r'\d+', response_clean)
+        if numbers:
+            n = int(numbers[0])
+            return 1 <= n <= 10
+        return False
+
+    elif verify == "contains_exact":
+        return check.lower() in response_lower
+
+    elif verify == "numbered_items":
+        items = re.findall(r'^\d+[\.\):]', response_clean, re.MULTILINE)
+        return len(items) >= check
+
+    elif verify == "starts_with":
+        return response_lower.startswith(check.lower())
+
+    elif verify == "starts_with_json":
+        stripped = response_clean.lstrip()
+        # May be in a code block
+        if "```" in stripped:
+            md_match = re.search(r'```(?:json)?\s*\n(.*?)```', stripped, re.DOTALL)
+            if md_match:
+                stripped = md_match.group(1).strip()
+        try:
+            obj = json.loads(stripped)
+            return "name" in obj and "age" in obj
+        except (json.JSONDecodeError, TypeError):
+            return stripped.startswith("{")
+
+    return False
+
+
+# ============================================================
+# MODEL RUNNER
+# ============================================================
+
+def load_model(model_path: str):
+    """Load model and tokenizer. Validates checkpoint integrity first."""
+    from transformers import AutoModelForImageTextToText, AutoTokenizer
+    from pathlib import Path
+
+    # Validate checkpoint before loading — catch stale/corrupt checkpoints
+    p = Path(model_path)
+    if p.is_dir():
+        config_file = p / "config.json"
+        if not config_file.exists():
+            raise FileNotFoundError(f"No config.json in {model_path} — stale or corrupt checkpoint!")
+        st_files = list(p.glob("*.safetensors"))
+        if not st_files and not list(p.glob("*.bin")):
+            raise FileNotFoundError(f"No model weights in {model_path} — checkpoint is empty!")
+        total_size = sum(f.stat().st_size for f in st_files) / 1e9 if st_files else 0
+        if st_files and total_size < 0.5:
+            raise ValueError(f"Model weights only {total_size:.2f} GB — likely corrupt checkpoint!")
+
+    print(f"  Loading model from {model_path}...")
+    try:
+        model = AutoModelForImageTextToText.from_pretrained(
+            model_path, dtype=torch.bfloat16,
+            device_map="auto", trust_remote_code=True
+        )
+    except (RuntimeError, torch.cuda.OutOfMemoryError) as e:
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        raise RuntimeError(f"Failed to load model (OOM or corrupt): {e}") from e
+
+    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+    return model, tokenizer
+
+
+def ask_model(model, tokenizer, question: str, max_tokens: int = 256,
+              disable_thinking: bool = True) -> str:
+    """Ask the model a question and get the response.
+
+    SPEED: disable_thinking=True adds /no_think to skip <think> blocks.
+    Qwen3-VL's thinking mode generates 100-500+ tokens of reasoning
+    before giving a 1-word answer. For benchmarks this wastes ~80% of
+    inference time. /no_think makes it answer directly.
+    """
+    # Build prompt — /no_think goes in the system message
+    if disable_thinking:
+        chat = (
+            "<|im_start|>system\n/no_think<|im_end|>\n"
+            f"<|im_start|>user\n{question}<|im_end|>\n"
+            "<|im_start|>assistant\n"
+        )
+    else:
+        chat = f"<|im_start|>user\n{question}<|im_end|>\n<|im_start|>assistant\n"
+
+    ids = tokenizer(chat, return_tensors="pt").to(model.device)
+
+    try:
+        with torch.no_grad():
+            out = model.generate(
+                **ids,
+                max_new_tokens=max_tokens,
+                do_sample=False,
+                temperature=None,
+                top_p=None,
+            )
+    except (torch.cuda.OutOfMemoryError, RuntimeError) as e:
+        # OOM during generation — clean up and return empty
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        print(f"    OOM during generation: {e}")
+        return "[OOM — no response generated]"
+
+    response = tokenizer.decode(out[0][ids["input_ids"].shape[1]:], skip_special_tokens=True)
+
+    # Guard against garbage output: if response is absurdly long, truncate
+    if len(response) > 5000:
+        response = response[:5000]
+
+    return response.strip()
+
+
+# ============================================================
+# RUN BENCHMARK
+# ============================================================
+
+def run_category(model, tokenizer, category_name: str, tests: list,
+                 check_fn, cfg: BenchmarkConfig) -> Dict:
+    """Run all tests in a category and return results."""
+    print(f"\n  === {category_name.upper()} ({len(tests)} questions) ===")
+
+    results = []
+    passed = 0
+    failed_questions = []
+
+    for i, test in enumerate(tests):
+        q = test["q"]
+        q_start = time.time()
+
+        try:
+            response = ask_model(model, tokenizer, q, cfg.max_new_tokens,
+                                 disable_thinking=cfg.disable_thinking)
+        except Exception as e:
+            print(f"    [ERROR] Question {i+1} crashed: {e}")
+            response = "[ERROR — question crashed]"
+            gc.collect()
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+
+        q_time = time.time() - q_start
+
+        try:
+            if check_fn == check_code_answer:
+                correct = check_fn(response, test)
+            elif check_fn == check_creativity_answer:
+                correct = check_fn(response, test)
+            elif check_fn == check_instruction_answer:
+                correct = check_fn(response, test)
+            elif check_fn == check_knowledge_answer:
+                correct = check_fn(response, test)
+            elif check_fn == check_reasoning_answer:
+                correct = check_fn(response, test.get("answer", ""))
+            else:
+                # Math
+                correct = check_fn(response, test.get("answer", ""))
+        except Exception as e:
+            print(f"    [ERROR] Answer checker crashed on question {i+1}: {e}")
+            correct = False
+
+        status = "PASS" if correct else "FAIL"
+        if correct:
+            passed += 1
+        else:
+            # Store the clean response (without thinking) for the failure report
+            clean_response = strip_thinking(response)
+            failed_questions.append({
+                "question": q,
+                "response": clean_response[:300],
+                "expected": test.get("answer", test.get("check", "N/A")),
+                "type": test.get("type", "unknown"),
+                "difficulty": test.get("difficulty", "unknown"),
+            })
+
+        # Print compact result
+        short_q = q[:60] + "..." if len(q) > 60 else q
+        print(f"    [{status}] {short_q} ({q_time:.1f}s)")
+        if not correct:
+            clean_r = strip_thinking(response)
+            short_r = clean_r[:80] + "..." if len(clean_r) > 80 else clean_r
+            print(f"          Got: {short_r}")
+
+        results.append({
+            "question": q,
+            "correct": correct,
+            "type": test.get("type", "unknown"),
+            "difficulty": test.get("difficulty", "unknown"),
+        })
+
+    total = len(tests)
+    # CRITICAL FIX: Score as 0.0-1.0 (NOT 0-100) so td_weakness.py targets work
+    score = passed / total if total > 0 else 0.0
+
+    print(f"\n    SCORE: {passed}/{total} ({score:.0%})")
+
+    # Per-type breakdown
+    type_scores = {}
+    for r in results:
+        t = r["type"]
+        if t not in type_scores:
+            type_scores[t] = {"correct": 0, "total": 0}
+        type_scores[t]["total"] += 1
+        if r["correct"]:
+            type_scores[t]["correct"] += 1
+
+    for t, s in sorted(type_scores.items()):
+        pct = s["correct"] / s["total"] if s["total"] > 0 else 0.0
+        print(f"      {t}: {s['correct']}/{s['total']} ({pct:.0%})")
+
+    return {
+        "category": category_name,
+        "score": score,  # 0.0-1.0 scale
+        "passed": passed,
+        "total": total,
+        "type_scores": {t: {"correct": s["correct"], "total": s["total"],
+                            "pct": s["correct"] / s["total"] if s["total"] > 0 else 0.0}
+                        for t, s in type_scores.items()},
+        "failed_questions": failed_questions,
+        "results": results,
+    }
+
+
+def run_benchmark(cfg: BenchmarkConfig = None) -> Dict:
+    """Run the full benchmark suite."""
+    if cfg is None:
+        cfg = BenchmarkConfig()
+
+    # Auto-detect model
+    if not cfg.model_path:
+        from pathlib import Path
+        base = Path("td_fuse_outputs/self_improve")
+        if base.exists():
+            for n in range(50, 0, -1):
+                d = base / f"improved_cycle{n}"
+                if d.exists() and list(d.glob("*.safetensors")):
+                    cfg.model_path = str(d)
+                    break
+        if not cfg.model_path:
+            healed = Path("td_fuse_outputs/reasoning_healed")
+            if healed.exists():
+                cfg.model_path = str(healed)
+        if not cfg.model_path:
+            raise FileNotFoundError("No model found!")
+
+    start = time.time()
+    print("=" * 60)
+    print(f"TD BENCHMARK SUITE v1")
+    print(f"Model: {cfg.model_path}")
+    print(f"Started: {time.strftime('%H:%M:%S')}")
+    print("=" * 60)
+
+    model, tokenizer = load_model(cfg.model_path)
+
+    # Run all categories
+    # SPEED: Per-category max tokens. Math/knowledge/reasoning answers are
+    # 1-10 tokens. Code needs 200+. Creativity needs 150+.
+    # Using 256 for a "yes/no" question wastes inference time.
+    CATEGORY_MAX_TOKENS = {
+        "math": 64,
+        "code": 300,  # Functions can be long
+        "reasoning": 64,
+        "creativity": 200,  # Stories/poems need room
+        "knowledge": 64,
+        "instruction_following": 128,  # JSON output, alphabet, etc.
+    }
+
+    categories = [
+        ("math", MATH_TESTS, check_math_answer),
+        ("code", CODE_TESTS, check_code_answer),
+        ("reasoning", REASONING_TESTS, check_reasoning_answer),
+        ("creativity", CREATIVITY_TESTS, check_creativity_answer),
+        ("knowledge", KNOWLEDGE_TESTS, check_knowledge_answer),
+        ("instruction_following", INSTRUCTION_TESTS, check_instruction_answer),
+    ]
+
+    all_results = {}
+    total_passed = 0
+    total_questions = 0
+
+    # Save original max_new_tokens and swap per-category
+    orig_max_tokens = cfg.max_new_tokens
+    for cat_name, tests, check_fn in categories:
+        cfg.max_new_tokens = CATEGORY_MAX_TOKENS.get(cat_name, orig_max_tokens)
+        result = run_category(model, tokenizer, cat_name, tests, check_fn, cfg)
+        all_results[cat_name] = result
+        total_passed += result["passed"]
+        total_questions += result["total"]
+    cfg.max_new_tokens = orig_max_tokens  # Restore
+
+    # Cleanup
+    del model
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+    elapsed = (time.time() - start) / 60
+    overall_score = total_passed / total_questions if total_questions > 0 else 0.0
+
+    # Print summary
+    print(f"\n{'='*60}")
+    print(f"BENCHMARK COMPLETE — {elapsed:.1f} min")
+    print(f"{'='*60}")
+    print(f"\n  REPORT CARD:")
+    print(f"  {'Category':<25} {'Score':>10} {'Grade':>8}")
+    print(f"  {'-'*45}")
+
+    for cat_name in ["math", "code", "reasoning", "creativity", "knowledge", "instruction_following"]:
+        r = all_results[cat_name]
+        score = r["score"]  # 0.0-1.0
+        if score >= 0.90: grade = "A"
+        elif score >= 0.80: grade = "B"
+        elif score >= 0.70: grade = "C"
+        elif score >= 0.60: grade = "D"
+        else: grade = "F"
+        bar = "█" * int(score * 20)
+        print(f"  {cat_name:<25} {r['passed']:>3}/{r['total']:<3} ({score:>5.0%})  {grade:>4}  {bar}")
+
+    print(f"\n  OVERALL: {total_passed}/{total_questions} ({overall_score:.0%})")
+    print(f"{'='*60}")
+
+    # Identify weaknesses (bottom 3 sub-categories)
+    all_types = []
+    for cat_name, result in all_results.items():
+        for type_name, type_score in result["type_scores"].items():
+            all_types.append({
+                "category": cat_name,
+                "type": type_name,
+                "pct": type_score["pct"],
+                "correct": type_score["correct"],
+                "total": type_score["total"],
+            })
+
+    weakest = sorted(all_types, key=lambda x: x["pct"])[:5]
+    print(f"\n  TOP 5 WEAKNESSES:")
+    for w in weakest:
+        print(f"    {w['category']}/{w['type']}: {w['correct']}/{w['total']} ({w['pct']:.0%})")
+
+    # Save report
+    out_dir = Path(cfg.output_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    report = {
+        "model_path": cfg.model_path,
+        "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
+        "duration_min": elapsed,
+        "overall_score": overall_score,
+        "total_passed": total_passed,
+        "total_questions": total_questions,
+        "categories": {name: {
+            "score": r["score"],
+            "passed": r["passed"],
+            "total": r["total"],
+            "type_scores": r["type_scores"],
+            "failed_questions": r["failed_questions"],
+        } for name, r in all_results.items()},
+        "weakest_areas": weakest,
+    }
+
+    # Support both output_path (exact path) and output_dir (auto-named)
+    if cfg.output_path:
+        report_path = Path(cfg.output_path)
+        report_path.parent.mkdir(parents=True, exist_ok=True)
+    else:
+        report_path = out_dir / f"benchmark_{time.strftime('%Y%m%d_%H%M%S')}.json"
+
+    try:
+        with open(report_path, "w") as f:
+            json.dump(report, f, indent=2)
+        print(f"\n  Report saved: {report_path}")
+    except OSError as e:
+        print(f"\n  WARNING: Could not save benchmark report: {e}")
+        print(f"  Results are still returned in memory — loop can continue.")
+
+    return report
+
+
+if __name__ == "__main__":
+    run_benchmark()

hugging/td_fuse/td_contamination.py

@@ -0,0 +1,193 @@
+"""
+TD Contamination Check — Benchmark Isolation Guard
+
+Ensures training data doesn't contain benchmark or holdout questions.
+If benchmark questions leak into training, scores become meaningless.
+
+This runs BEFORE training starts. If contamination is found, it:
+  1. Removes the contaminated samples from training data
+  2. Logs what was removed and why
+  3. Returns the cleaned data
+
+Usage:
+  from td_contamination import check_contamination, get_all_benchmark_questions
+  clean_texts = check_contamination(train_texts)
+
+  # Or standalone:
+  python td_contamination.py path/to/training_data.jsonl
+"""
+
+import json
+import re
+from pathlib import Path
+from typing import List, Set, Tuple
+
+
+def get_all_benchmark_questions() -> Set[str]:
+    """
+    Collect ALL questions from td_benchmark.py and td_holdout.py.
+    Returns a set of normalized question strings.
+    """
+    questions = set()
+
+    # Import benchmark questions
+    try:
+        from td_benchmark import (MATH_TESTS, CODE_TESTS, REASONING_TESTS,
+                                   CREATIVITY_TESTS, KNOWLEDGE_TESTS, INSTRUCTION_TESTS)
+        for bank in [MATH_TESTS, CODE_TESTS, REASONING_TESTS,
+                     CREATIVITY_TESTS, KNOWLEDGE_TESTS, INSTRUCTION_TESTS]:
+            for test in bank:
+                questions.add(_normalize(test["q"]))
+    except ImportError:
+        print("  WARNING: Could not import td_benchmark — skipping main benchmark")
+
+    # Import holdout questions
+    try:
+        from td_holdout import (HOLDOUT_MATH, HOLDOUT_CODE, HOLDOUT_REASONING,
+                                 HOLDOUT_CREATIVITY, HOLDOUT_KNOWLEDGE, HOLDOUT_INSTRUCTION)
+        for bank in [HOLDOUT_MATH, HOLDOUT_CODE, HOLDOUT_REASONING,
+                     HOLDOUT_CREATIVITY, HOLDOUT_KNOWLEDGE, HOLDOUT_INSTRUCTION]:
+            for test in bank:
+                questions.add(_normalize(test["q"]))
+    except ImportError:
+        print("  WARNING: Could not import td_holdout — skipping holdout benchmark")
+
+    return questions
+
+
+def _normalize(text: str) -> str:
+    """Normalize text for fuzzy matching."""
+    text = text.lower().strip()
+    # Remove punctuation except basic math operators
+    text = re.sub(r'[^\w\s+\-*/=^]', '', text)
+    # Collapse whitespace
+    text = re.sub(r'\s+', ' ', text)
+    return text
+
+
+def _extract_question_from_chat(text: str) -> str:
+    """Extract the user question from a chat-formatted training sample."""
+    # Match <|im_start|>user\n...<|im_end|>
+    match = re.search(r'<\|im_start\|>user\n(.*?)<\|im_end\|>', text, re.DOTALL)
+    if match:
+        return match.group(1).strip()
+    return text
+
+
+def check_contamination(train_texts: List[str], verbose: bool = True) -> List[str]:
+    """
+    Check training data for benchmark question contamination.
+
+    Args:
+        train_texts: List of training text strings
+        verbose: Print details about contamination found
+
+    Returns:
+        Cleaned list with contaminated samples removed
+    """
+    benchmark_questions = get_all_benchmark_questions()
+    if not benchmark_questions:
+        if verbose:
+            print("  Contamination check: No benchmark questions loaded — SKIP")
+        return train_texts
+
+    if verbose:
+        print(f"\n  Contamination check: scanning {len(train_texts)} samples against "
+              f"{len(benchmark_questions)} benchmark questions...")
+
+    clean = []
+    removed = []
+    # Also build n-gram set for partial matching (catches paraphrased questions)
+    benchmark_ngrams = set()
+    for q in benchmark_questions:
+        words = q.split()
+        # 5-grams from each question
+        for i in range(max(0, len(words) - 4)):
+            benchmark_ngrams.add(' '.join(words[i:i+5]))
+
+    for text in train_texts:
+        question = _extract_question_from_chat(text)
+        norm_q = _normalize(question)
+
+        # Check 1: Exact match
+        if norm_q in benchmark_questions:
+            removed.append(("exact", question[:80]))
+            continue
+
+        # Check 2: Substring match — benchmark question appears inside training text
+        contaminated = False
+        for bq in benchmark_questions:
+            if len(bq) > 15 and bq in _normalize(text):
+                removed.append(("substring", f"'{bq[:50]}' found in sample"))
+                contaminated = True
+                break
+
+        if contaminated:
+            continue
+
+        # Check 3: High n-gram overlap (catches paraphrases)
+        norm_words = norm_q.split()
+        if len(norm_words) >= 5:
+            sample_ngrams = set()
+            for i in range(len(norm_words) - 4):
+                sample_ngrams.add(' '.join(norm_words[i:i+5]))
+            if sample_ngrams and benchmark_ngrams:
+                overlap = len(sample_ngrams & benchmark_ngrams) / max(len(sample_ngrams), 1)
+                if overlap > 0.5:  # More than half the 5-grams match a benchmark question
+                    removed.append(("ngram_overlap", f"{overlap:.0%} overlap: {question[:60]}"))
+                    continue
+
+        clean.append(text)
+
+    if verbose:
+        if removed:
+            print(f"  CONTAMINATION FOUND: {len(removed)} samples removed!")
+            for reason, detail in removed[:10]:  # Show first 10
+                print(f"    [{reason}] {detail}")
+            if len(removed) > 10:
+                print(f"    ... and {len(removed) - 10} more")
+        else:
+            print(f"  Contamination check: CLEAN (0 matches found)")
+
+    return clean
+
+
+def check_file(filepath: str) -> Tuple[int, int]:
+    """Check a JSONL file for contamination. Returns (clean_count, removed_count)."""
+    texts = []
+    with open(filepath) as f:
+        for line in f:
+            line = line.strip()
+            if not line:
+                continue
+            try:
+                item = json.loads(line)
+                text = item.get("text", "")
+                if text:
+                    texts.append(text)
+            except json.JSONDecodeError:
+                continue
+
+    clean = check_contamination(texts)
+    return len(clean), len(texts) - len(clean)
+
+
+if __name__ == "__main__":
+    import sys
+
+    if len(sys.argv) < 2:
+        # Just report benchmark question count
+        questions = get_all_benchmark_questions()
+        print(f"Loaded {len(questions)} benchmark questions to check against")
+        print("\nUsage: python td_contamination.py path/to/training_data.jsonl")
+        sys.exit(0)
+
+    filepath = sys.argv[1]
+    if not Path(filepath).exists():
+        print(f"File not found: {filepath}")
+        sys.exit(1)
+
+    clean_count, removed_count = check_file(filepath)
+    print(f"\nResult: {clean_count} clean, {removed_count} contaminated")
+    if removed_count > 0:
+        print("ACTION: Training on this data would inflate benchmark scores!")

hugging/td_fuse/td_data_finder.py

@@ -0,0 +1,615 @@
+"""
+TD Dataset Finder v1 — Search Everywhere for Training Data
+
+Takes the weakness report from td_weakness.py and finds training data
+to fix those weaknesses. Searches multiple sources in priority order.
+
+This is Step 3 of the self-improvement loop.
+
+Sources (in order of priority):
+  1. HuggingFace Datasets API — biggest, easiest to download
+  2. GitHub — repos with training data, benchmarks, problem sets
+  3. Kaggle — competitions and datasets
+  4. (Future) Web scraping — Stack Overflow, math forums, etc.
+
+Input:
+  - WeaknessReport from td_weakness.py
+
+Output:
+  - Formatted training data in chat template, ready for training
+  - Data allocation follows the weakness report (worse = more data)
+"""
+
+import json
+import os
+import random
+import time
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple
+from dataclasses import dataclass, field
+
+
+# ============================================================
+# DATA STRUCTURES
+# ============================================================
+
+@dataclass
+class DatasetResult:
+    """A dataset found from searching."""
+    name: str
+    source: str  # "huggingface", "github", "kaggle", "web"
+    url: str
+    description: str
+    size_estimate: int  # approximate number of samples
+    relevance_score: float  # 0-1, how well it matches the weakness
+    category: str  # which weakness category it's for
+
+
+@dataclass
+class FormattedSample:
+    """A single training sample in chat template format."""
+    system: str
+    user: str
+    assistant: str
+    category: str
+    source: str
+
+
+@dataclass
+class TrainingData:
+    """Complete training data package ready for the trainer."""
+    samples: List[FormattedSample]
+    sources_used: List[DatasetResult]
+    allocation: Dict[str, int]  # category -> actual samples collected
+    total_samples: int
+    categories_covered: List[str]
+    categories_missing: List[str]  # categories we couldn't find data for
+
+
+# ============================================================
+# CHAT TEMPLATE FORMATTER
+# ============================================================
+
+def format_chat_template(
+    system: str,
+    user: str,
+    assistant: str,
+) -> str:
+    """
+    Format a sample into Qwen3-VL's chat template.
+    This is the <|im_start|> format that selfimprove.py already uses.
+    """
+    parts = []
+    if system:
+        parts.append(f"<|im_start|>system\n{system}<|im_end|>")
+    parts.append(f"<|im_start|>user\n{user}<|im_end|>")
+    parts.append(f"<|im_start|>assistant\n{assistant}<|im_end|>")
+    return "\n".join(parts)
+
+
+def format_sample_to_dict(sample: FormattedSample) -> Dict:
+    """Convert a FormattedSample to the dict format selfimprove.py expects."""
+    return {
+        "text": format_chat_template(sample.system, sample.user, sample.assistant),
+        "category": sample.category,
+        "source": sample.source,
+    }
+
+
+# ============================================================
+# HUGGINGFACE SEARCH
+# ============================================================
+
+def search_huggingface(
+    keywords: List[str],
+    tags: List[str],
+    category: str,
+    max_datasets: int = 5,
+) -> List[DatasetResult]:
+    """
+    Search HuggingFace Datasets API for matching datasets.
+
+    Uses the huggingface_hub library to search by keyword and tags.
+    Returns metadata about found datasets (doesn't download yet).
+    """
+    results = []
+
+    try:
+        from huggingface_hub import HfApi
+        api = HfApi()
+
+        # Search by each keyword
+        seen = set()
+        for keyword in keywords[:3]:  # Top 3 keywords
+            try:
+                datasets = api.list_datasets(
+                    search=keyword,
+                    sort="downloads",
+                    direction=-1,
+                    limit=max_datasets,
+                )
+
+                for ds in datasets:
+                    if ds.id in seen:
+                        continue
+                    seen.add(ds.id)
+
+                    # Estimate relevance based on tag overlap
+                    ds_tags = set(ds.tags or [])
+                    tag_overlap = len(ds_tags.intersection(set(tags)))
+                    relevance = min(0.5 + (tag_overlap * 0.15), 1.0)
+
+                    # Boost if keyword appears in name
+                    if any(kw.lower() in ds.id.lower() for kw in keyword.split()):
+                        relevance = min(relevance + 0.2, 1.0)
+
+                    results.append(DatasetResult(
+                        name=ds.id,
+                        source="huggingface",
+                        url=f"https://huggingface.co/datasets/{ds.id}",
+                        description=ds.id,  # HF API doesn't always return description
+                        size_estimate=ds.downloads or 0,  # Use downloads as proxy
+                        relevance_score=relevance,
+                        category=category,
+                    ))
+            except Exception as e:
+                print(f"    Warning: HF search failed for '{keyword}': {e}")
+                continue
+
+    except ImportError:
+        print("    Warning: huggingface_hub not installed. Install with: pip install huggingface_hub")
+        return []
+
+    # Sort by relevance
+    results.sort(key=lambda r: r.relevance_score, reverse=True)
+    return results[:max_datasets]
+
+
+def download_huggingface_dataset(
+    dataset_name: str,
+    category: str,
+    num_samples: int,
+    split: str = "train",
+) -> List[FormattedSample]:
+    """
+    Download and format samples from a HuggingFace dataset.
+
+    Handles common dataset formats:
+    - QA pairs (question/answer columns)
+    - Instruction pairs (instruction/output columns)
+    - Chat format (conversations column)
+    - Text only (text column) — wrap in QA format
+    """
+    samples = []
+
+    try:
+        from datasets import load_dataset
+
+        print(f"    Downloading {dataset_name} (split={split})...")
+        try:
+            ds = load_dataset(dataset_name, split=split, streaming=True)
+        except Exception:
+            # Try without split specification
+            try:
+                ds = load_dataset(dataset_name, streaming=True)
+                # Get first available split
+                if hasattr(ds, 'keys'):
+                    split_name = list(ds.keys())[0]
+                    ds = ds[split_name]
+                else:
+                    return []
+            except Exception as e:
+                print(f"    Failed to load {dataset_name}: {e}")
+                return []
+
+        # Collect samples (streaming so we don't download everything)
+        # Timeout: 5 min per dataset to prevent hanging on slow/broken sources
+        import time as _time
+        download_start = _time.time()
+        download_timeout = 300  # 5 minutes max per dataset
+        count = 0
+        stall_count = 0
+        for item in ds:
+            if count >= num_samples:
+                break
+            if _time.time() - download_start > download_timeout:
+                print(f"    Timeout: {dataset_name} took >{download_timeout}s — moving on with {count} samples")
+                break
+
+            sample = _format_hf_item(item, category, dataset_name)
+            if sample:
+                samples.append(sample)
+                count += 1
+                stall_count = 0
+            else:
+                stall_count += 1
+                # If 500 consecutive items produce nothing, dataset format is wrong
+                if stall_count > 500:
+                    print(f"    {dataset_name}: 500 consecutive unparseable items — skipping")
+                    break
+
+        print(f"    Got {len(samples)} samples from {dataset_name}")
+
+    except ImportError:
+        print("    Warning: datasets library not installed. Install with: pip install datasets")
+    except Exception as e:
+        print(f"    Error downloading {dataset_name}: {e}")
+
+    return samples
+
+
+def _format_hf_item(item: Dict, category: str, source: str) -> Optional[FormattedSample]:
+    """
+    Try to format a HuggingFace dataset item into a chat sample.
+    Handles many common column naming conventions.
+    """
+    # Guard: item must be a dict
+    if not isinstance(item, dict):
+        return None
+
+    # Common column name patterns for user/assistant content
+    user_keys = ["question", "instruction", "input", "prompt", "query", "problem", "text"]
+    assistant_keys = ["answer", "output", "response", "solution", "target", "completion"]
+    system_keys = ["system", "system_prompt", "context"]
+
+    user_text = None
+    assistant_text = None
+    system_text = ""
+
+    # Try to find user content (guard: value might be non-string like int, list, etc.)
+    for key in user_keys:
+        if key in item and item[key]:
+            val = item[key]
+            if isinstance(val, (str, int, float)):
+                user_text = str(val).strip()
+                if user_text:
+                    break
+            # Skip lists, dicts, None-like values
+
+    # Try to find assistant content
+    for key in assistant_keys:
+        if key in item and item[key]:
+            val = item[key]
+            if isinstance(val, (str, int, float)):
+                assistant_text = str(val).strip()
+                if assistant_text:
+                    break
+
+    # Try to find system content
+    for key in system_keys:
+        if key in item and item[key]:
+            val = item[key]
+            if isinstance(val, (str, int, float)):
+                system_text = str(val).strip()
+                break
+
+    # Handle conversation format
+    if not user_text and "conversations" in item:
+        convs = item["conversations"]
+        if isinstance(convs, list) and len(convs) >= 2:
+            for c in convs:
+                if not isinstance(c, dict):
+                    continue  # Skip malformed conversation entries
+                role = str(c.get("role", c.get("from", ""))).strip()
+                content = c.get("content", c.get("value", ""))
+                if content is None:
+                    continue
+                content = str(content).strip()
+                if role in ["user", "human"] and not user_text:
+                    user_text = content
+                elif role in ["assistant", "gpt", "bot"] and not assistant_text:
+                    assistant_text = content
+                elif role == "system" and not system_text:
+                    system_text = content
+
+    # Handle messages format (OpenAI-style)
+    if not user_text and "messages" in item:
+        msgs = item["messages"]
+        if isinstance(msgs, list):
+            for m in msgs:
+                if not isinstance(m, dict):
+                    continue  # Skip malformed message entries
+                role = str(m.get("role", "")).strip()
+                content = m.get("content", "")
+                if content is None:
+                    continue
+                content = str(content).strip()
+                if role == "user" and not user_text:
+                    user_text = content
+                elif role == "assistant" and not assistant_text:
+                    assistant_text = content
+                elif role == "system" and not system_text:
+                    system_text = content
+
+    # Handle MMLU-style multiple choice: question + choices + answer_index
+    if not assistant_text and "choices" in item and "answer" in item:
+        choices = item["choices"]
+        answer_idx = item["answer"]
+        if isinstance(choices, list) and len(choices) > 0 and len(choices) <= 26:
+            # Format as multiple choice question (cap at 26 = A-Z)
+            if user_text:
+                choice_strs = [f"{chr(65+i)}. {c}" for i, c in enumerate(choices)]
+                user_text = user_text + "\n" + "\n".join(choice_strs)
+            if isinstance(answer_idx, int) and 0 <= answer_idx < len(choices):
+                assistant_text = f"{chr(65+answer_idx)}. {choices[answer_idx]}"
+            elif isinstance(answer_idx, str) and len(answer_idx) == 1 and answer_idx.isalpha():
+                idx = ord(answer_idx.upper()) - 65
+                if 0 <= idx < len(choices):
+                    assistant_text = f"{answer_idx.upper()}. {choices[idx]}"
+
+    # Handle GSM8K-style: answer contains "#### final_number"
+    if assistant_text and "####" in assistant_text:
+        # Keep the full chain-of-thought — it's good training data for reasoning
+        pass
+
+    # Need at least user + assistant
+    if not user_text or not assistant_text:
+        return None
+
+    # Skip very short or very long samples
+    if len(user_text) < 10 or len(assistant_text) < 5:  # Allow short answers (like "42")
+        return None
+    if len(user_text) > 5000 or len(assistant_text) > 10000:
+        return None
+
+    return FormattedSample(
+        system=system_text,
+        user=user_text,
+        assistant=assistant_text,
+        category=category,
+        source=source,
+    )
+
+
+# ============================================================
+# WELL-KNOWN DATASETS PER CATEGORY
+# ============================================================
+
+# If the search finds nothing, fall back to these known-good datasets
+# Verified dataset names on HuggingFace (March 2026)
+# Each tuple: (dataset_name, config_or_subset)
+# If a dataset fails to load, it moves to the next one silently
+FALLBACK_DATASETS = {
+    "math": [
+        ("openai/gsm8k", "main"),          # 8.8K train, question/answer with ####
+        ("lighteval/MATH", "all"),           # 12.5K, competition math
+        ("microsoft/orca-math-word-problems-200k", None),  # 200K word problems
+    ],
+    "code": [
+        ("openai/openai_humaneval", None),   # 164 problems — small but high quality
+        ("google-research-datasets/mbpp", None),  # FIX: full name (short "mbpp" fails)
+        ("codeparrot/apps", "introductory"), # Competitive programming
+    ],
+    "reasoning": [
+        ("allenai/ai2_arc", "ARC-Challenge"),  # 2.6K challenge set
+        ("Rowan/hellaswag", None),              # 40K commonsense
+        ("allenai/openbookqa", "main"),         # FIX: replaced openai_prm800k (not a public dataset)
+    ],
+    "creativity": [
+        ("teknium/OpenHermes-2.5", None),   # 1M instruction pairs
+        ("Open-Orca/OpenOrca", None),        # 4.2M conversations
+    ],
+    "knowledge": [
+        ("TIGER-Lab/MMLU-Pro", None),        # MMLU harder version
+        ("cais/mmlu", "all"),                 # Classic MMLU
+        ("truthfulqa/truthful_qa", "generation"),  # 817 truthfulness questions
+    ],
+    "instruction_following": [
+        ("HuggingFaceH4/ultrafeedback_binarized", None),  # 61K preference pairs
+        ("tatsu-lab/alpaca", None),           # 52K instructions
+    ],
+}
+
+
+# ============================================================
+# MAIN DATA FINDER
+# ============================================================
+
+def find_training_data(
+    search_plan: List[Dict],
+    data_allocation: Dict[str, int],
+    output_dir: str = "td_fuse_outputs/training_data",
+    hf_token: Optional[str] = None,
+) -> TrainingData:
+    """
+    Search for and download training data based on the weakness report.
+
+    This is the main entry point that td_loop.py calls.
+
+    Args:
+        search_plan: From WeaknessReport.search_plan
+        data_allocation: From WeaknessReport.data_allocation
+        output_dir: Where to save the formatted data
+        hf_token: HuggingFace token for private datasets
+
+    Returns:
+        TrainingData with formatted samples ready for training
+    """
+    print("\n" + "=" * 60)
+    print("TD DATASET FINDER v1")
+    print("=" * 60)
+
+    if hf_token:
+        os.environ["HF_TOKEN"] = hf_token
+
+    all_samples = []
+    all_sources = []
+    categories_covered = []
+    categories_missing = []
+
+    for plan_item in search_plan:
+        category = plan_item["category"]
+        needed = data_allocation.get(category, 1000)
+        keywords = plan_item["keywords"]
+        hf_tags = plan_item.get("hf_tags", [])
+
+        print(f"\n--- Searching for: {category} ({needed} samples needed) ---")
+
+        # Step 1: Search HuggingFace
+        print(f"  Searching HuggingFace...")
+        hf_results = search_huggingface(keywords, hf_tags, category)
+
+        if hf_results:
+            print(f"  Found {len(hf_results)} datasets on HuggingFace:")
+            for r in hf_results[:3]:
+                print(f"    - {r.name} (relevance: {r.relevance_score:.2f})")
+
+        # Step 2: Download from best matches
+        category_samples = []
+        for result in hf_results:
+            if len(category_samples) >= needed:
+                break
+
+            remaining = needed - len(category_samples)
+            samples = download_huggingface_dataset(
+                dataset_name=result.name,
+                category=category,
+                num_samples=remaining,
+            )
+            category_samples.extend(samples)
+            if samples:
+                all_sources.append(result)
+
+        # Step 3: If not enough, try fallback datasets
+        if len(category_samples) < needed:
+            print(f"  Only got {len(category_samples)}/{needed}. Trying fallback datasets...")
+            fallbacks = FALLBACK_DATASETS.get(category, [])
+            try:
+                from datasets import load_dataset  # Import once outside loop
+            except ImportError:
+                print("    Warning: datasets library not installed — cannot try fallbacks")
+                fallbacks = []  # Skip the fallback loop
+            for ds_name, config in fallbacks:
+                if len(category_samples) >= needed:
+                    break
+
+                remaining = needed - len(category_samples)
+                try:
+                    if config:
+                        ds = load_dataset(ds_name, config, streaming=True, split="train")
+                    else:
+                        ds = load_dataset(ds_name, streaming=True, split="train")
+
+                    count = 0
+                    for item in ds:
+                        if count >= remaining:
+                            break
+                        sample = _format_hf_item(item, category, ds_name)
+                        if sample:
+                            category_samples.append(sample)
+                            count += 1
+
+                    if count > 0:
+                        print(f"    Got {count} from fallback: {ds_name}")
+                        all_sources.append(DatasetResult(
+                            name=ds_name,
+                            source="huggingface_fallback",
+                            url=f"https://huggingface.co/datasets/{ds_name}",
+                            description=f"Fallback dataset for {category}",
+                            size_estimate=count,
+                            relevance_score=0.5,
+                            category=category,
+                        ))
+                except Exception as e:
+                    print(f"    Fallback {ds_name} failed: {e}")
+                    continue
+
+        # Track what we got
+        if category_samples:
+            all_samples.extend(category_samples)
+            categories_covered.append(category)
+            print(f"  Total for {category}: {len(category_samples)} samples")
+        else:
+            categories_missing.append(category)
+            print(f"  WARNING: No data found for {category}!")
+
+    # Shuffle all samples together
+    random.shuffle(all_samples)
+
+    # Save to disk
+    output_path = Path(output_dir)
+    try:
+        output_path.mkdir(parents=True, exist_ok=True)
+    except OSError as e:
+        print(f"\n  ERROR: Cannot create output dir {output_dir}: {e}")
+        # Still return what we collected even if can't save
+        actual_allocation = {}
+        for s in all_samples:
+            actual_allocation[s.category] = actual_allocation.get(s.category, 0) + 1
+        return TrainingData(
+            samples=all_samples, sources_used=all_sources,
+            allocation=actual_allocation, total_samples=len(all_samples),
+            categories_covered=categories_covered, categories_missing=categories_missing,
+        )
+
+    # Save as JSONL (one sample per line)
+    data_file = output_path / "training_data.jsonl"
+    try:
+        with open(data_file, "w") as f:
+            for sample in all_samples:
+                f.write(json.dumps(format_sample_to_dict(sample)) + "\n")
+        print(f"\nSaved {len(all_samples)} samples to {data_file}")
+    except OSError as e:
+        print(f"\n  WARNING: Failed to save training data: {e} (disk may be full)")
+
+    # Save metadata
+    meta_file = output_path / "data_metadata.json"
+    actual_allocation = {}
+    for s in all_samples:
+        actual_allocation[s.category] = actual_allocation.get(s.category, 0) + 1
+
+    meta = {
+        "total_samples": len(all_samples),
+        "allocation": actual_allocation,
+        "sources": [
+            {"name": s.name, "source": s.source, "category": s.category, "url": s.url}
+            for s in all_sources
+        ],
+        "categories_covered": categories_covered,
+        "categories_missing": categories_missing,
+    }
+    try:
+        with open(meta_file, "w") as f:
+            json.dump(meta, f, indent=2)
+        print(f"Saved metadata to {meta_file}")
+    except OSError as e:
+        print(f"  WARNING: Failed to save metadata: {e}")
+
+    return TrainingData(
+        samples=all_samples,
+        sources_used=all_sources,
+        allocation=actual_allocation,
+        total_samples=len(all_samples),
+        categories_covered=categories_covered,
+        categories_missing=categories_missing,
+    )
+
+
+# ============================================================
+# STANDALONE USAGE
+# ============================================================
+
+if __name__ == "__main__":
+    import sys
+
+    if len(sys.argv) < 2:
+        print("Usage: python td_data_finder.py <weakness_report.json> [output_dir]")
+        print("\nSearches for training data based on the weakness report.")
+        sys.exit(1)
+
+    # Load weakness report
+    from td_weakness import load_report
+    report = load_report(sys.argv[1])
+
+    output_dir = sys.argv[2] if len(sys.argv) > 2 else "td_fuse_outputs/training_data"
+
+    # Find data
+    training_data = find_training_data(
+        search_plan=report.search_plan,
+        data_allocation=report.data_allocation,
+        output_dir=output_dir,
+    )
+
+    print(f"\nDone! Got {training_data.total_samples} samples across {len(training_data.categories_covered)} categories.")
+    if training_data.categories_missing:
+        print(f"Missing data for: {', '.join(training_data.categories_missing)}")

hugging/td_fuse/td_holdout.py

@@ -0,0 +1,426 @@
+"""
+TD Holdout Benchmark — The Truth Test
+
+This is a SEPARATE benchmark that NEVER influences training decisions.
+It only runs at the start and end to verify improvement is real.
+
+Why this exists:
+  The main benchmark (td_benchmark.py) has 85 fixed questions.
+  After enough cycles, the model might get good at THOSE specific
+  questions without actually getting smarter. This is overfitting.
+
+  The holdout uses completely different questions testing the same
+  skills. If the main benchmark shows improvement but the holdout
+  doesn't, we know we're overfitting.
+
+Rules:
+  1. NEVER use holdout results to decide training data
+  2. NEVER share holdout questions with the training pipeline
+  3. Only run holdout at: first cycle (baseline) and final cycle (verify)
+  4. If holdout score drops while main benchmark rises = OVERFITTING
+
+Usage:
+  python td_holdout.py                    # Run on latest model
+  python td_holdout.py path/to/model      # Run on specific model
+"""
+
+import json
+import time
+import gc
+import re
+from pathlib import Path
+from typing import Dict
+
+
+# ============================================================
+# HOLDOUT QUESTIONS — completely different from td_benchmark.py
+# ============================================================
+
+HOLDOUT_MATH = [
+    # Different numbers, same skills as main benchmark
+    {"q": "What is 63 * 17?", "answer": "1071", "type": "arithmetic"},
+    {"q": "What is 892 - 456?", "answer": "436", "type": "arithmetic"},
+    {"q": "What is 25% of 360?", "answer": "90", "type": "arithmetic"},
+    {"q": "Solve for x: 5x - 3 = 27", "answer": "6", "type": "algebra"},
+    {"q": "If f(x) = 3x - 1, what is f(5)?", "answer": "14", "type": "algebra"},
+    {"q": "A book costs $15. If there's a 20% discount, what's the price?", "answer": "12", "type": "word_problem"},
+    {"q": "A tank fills at 3 liters per minute. How long to fill a 45 liter tank?", "answer": "15", "type": "word_problem"},
+    {"q": "What is the sum of the first 5 prime numbers?", "answer": "28", "type": "number_theory"},
+    {"q": "What is 3^5?", "answer": "243", "type": "exponents"},
+    {"q": "What is the least common multiple of 6 and 8?", "answer": "24", "type": "number_theory"},
+]
+
+HOLDOUT_CODE = [
+    {"q": "Write a Python function called 'sum_list' that takes a list of numbers and returns their sum. Just the function.",
+     "check": "sum_list", "verify": "callable",
+     "tests": [{"args": [[1,2,3]], "expected": 6}, {"args": [[]], "expected": 0}]},
+    {"q": "Write a Python function called 'remove_duplicates' that takes a list and returns a new list with duplicates removed, preserving order. Just the function.",
+     "check": "remove_duplicates", "verify": "callable",
+     "tests": [{"args": [[1,2,2,3,1]], "expected": [1,2,3]}]},
+    {"q": "Write a Python function called 'is_prime' that returns True if a number is prime. Just the function.",
+     "check": "is_prime", "verify": "callable",
+     "tests": [{"args": [7], "expected": True}, {"args": [4], "expected": False}, {"args": [1], "expected": False}]},
+    {"q": "Write a Python function called 'count_words' that counts the number of words in a string. Just the function.",
+     "check": "count_words", "verify": "callable",
+     "tests": [{"args": ["hello world foo"], "expected": 3}, {"args": [""], "expected": 0}]},
+    {"q": "This code has a bug. Fix it:\ndef average(numbers):\n    return sum(numbers) / len(numbers)",
+     "check": "len(numbers) == 0", "verify": "contains_any",
+     "alt_checks": ["not numbers", "len(numbers) == 0", "if numbers", "if not", "len(numbers) < 1", "ZeroDivisionError"]},
+]
+
+HOLDOUT_REASONING = [
+    {"q": "If all dogs are mammals and some mammals can swim, can all dogs swim? Answer yes or no.", "answer": "no", "type": "syllogism"},
+    {"q": "What comes next: 5, 10, 20, 40, ?", "answer": "80", "type": "pattern"},
+    {"q": "What comes next: 1, 3, 6, 10, 15, ?", "answer": "21", "type": "pattern"},
+    {"q": "There are 5 people in a room. Each shakes hands with everyone else exactly once. How many handshakes total?", "answer": "10", "type": "multi_step"},
+    {"q": "A farmer has chickens and cows. He counts 20 heads and 56 legs. How many chickens does he have?", "answer": "12", "type": "multi_step"},
+    {"q": "If you rearrange the letters 'CIFAIPC' you get the name of a what? Ocean or country?", "answer": "ocean", "type": "common_sense"},
+    {"q": "Is it possible for your mother's brother to be a woman? Answer yes or no.", "answer": "no", "type": "common_sense"},
+]
+
+HOLDOUT_CREATIVITY = [
+    {"q": "Write a 4-line poem about rain.", "min_lines": 3, "min_words": 15, "type": "poem"},
+    {"q": "Explain gravity to a 5 year old in 2-3 sentences.", "min_lines": 1, "min_words": 15, "type": "explain_simple"},
+    {"q": "Come up with 3 creative names for a pet goldfish.", "min_items": 3, "type": "brainstorm"},
+    {"q": "Rewrite this to be more dramatic: 'The sun set behind the mountains.'",
+     "min_words": 5, "type": "rewrite", "must_not_contain": "The sun set behind the mountains"},
+]
+
+HOLDOUT_KNOWLEDGE = [
+    {"q": "What is the largest planet in our solar system?", "answer": "jupiter", "type": "science"},
+    {"q": "What element has the atomic number 1?", "answer": "hydrogen", "type": "science"},
+    {"q": "Who wrote Romeo and Juliet?", "answer": "shakespeare", "type": "history", "alt_answers": ["william shakespeare"]},
+    {"q": "What is the capital of Australia?", "answer": "canberra", "type": "geography"},
+    {"q": "What year did the first iPhone come out?", "answer": "2007", "type": "tech"},
+    {"q": "What does RAM stand for?", "answer": "random access memory", "type": "tech"},
+]
+
+HOLDOUT_INSTRUCTION = [
+    {"q": "Respond with only the word 'goodbye'. Nothing else.",
+     "check": "goodbye", "verify": "exact_word", "type": "exact"},
+    {"q": "List exactly 4 animals, one per line. Nothing else.",
+     "check": 4, "verify": "line_count", "type": "format"},
+    {"q": "Write a sentence that contains exactly 7 words.",
+     "check": 7, "verify": "word_count", "type": "format"},
+    {"q": "Give me a number between 50 and 100. Just the number.",
+     "check": "number_50_100", "verify": "number_range_custom", "type": "exact"},
+    {"q": "Say 'no' if 10 > 20, otherwise say 'yes'. Only one word.",
+     "check": "yes", "verify": "exact_word", "type": "exact"},
+]
+
+
+# ============================================================
+# ANSWER CHECKING (reuses td_benchmark logic)
+# ============================================================
+
+def strip_thinking(response: str) -> str:
+    cleaned = re.sub(r'<think>.*?</think>', '', response, flags=re.DOTALL)
+    cleaned = re.sub(r'<think>.*$', '', cleaned, flags=re.DOTALL)
+    return cleaned.strip()
+
+
+def check_math(response: str, expected: str) -> bool:
+    response_clean = strip_thinking(response).lower().replace(",", "")
+    numbers = re.findall(r'-?\d+\.?\d*', response_clean)
+    for n in numbers:
+        try:
+            if float(n) == float(expected):
+                return True
+        except ValueError:
+            continue
+    return False
+
+
+def check_code(response: str, test: dict) -> bool:
+    response = strip_thinking(response)
+    code_block = response
+    md_match = re.search(r'```(?:python)?\s*\n(.*?)```', response, re.DOTALL)
+    if md_match:
+        code_block = md_match.group(1)
+    elif "def " in code_block:
+        def_start = code_block.index("def ")
+        code_block = code_block[def_start:]
+
+    if test["verify"] == "callable":
+        try:
+            import signal
+
+            def _timeout(signum, frame):
+                raise TimeoutError()
+
+            old = signal.signal(signal.SIGALRM, _timeout)
+            signal.alarm(10)
+            try:
+                ns = {}
+                exec(code_block.strip(), ns)
+                func = ns.get(test["check"])
+                if not func or not callable(func):
+                    return False
+                for t in test.get("tests", []):
+                    result = func(*t["args"])
+                    if result != t["expected"]:
+                        return False
+                return True
+            finally:
+                signal.alarm(0)
+                signal.signal(signal.SIGALRM, old)
+        except Exception:
+            return False
+    elif test["verify"] == "contains_any":
+        checks = test.get("alt_checks", [test["check"]])
+        return any(c.lower() in response.lower() for c in checks)
+    return False
+
+
+def check_reasoning(response: str, expected: str) -> bool:
+    response_lower = strip_thinking(response).lower().strip()
+    if expected.lower() in response_lower:
+        return True
+    if expected.lower() in ["yes", "no"]:
+        first = response_lower.split()[0].strip(".,!?") if response_lower.split() else ""
+        return first == expected.lower()
+    numbers = re.findall(r'-?\d+\.?\d*', response_lower)
+    return expected.lower() in numbers
+
+
+def check_creativity(response: str, test: dict) -> bool:
+    response = strip_thinking(response)
+    lines = [l.strip() for l in response.strip().split("\n") if l.strip()]
+    words = response.split()
+    if "must_not_contain" in test and test["must_not_contain"].lower() in response.lower():
+        return False
+    if "min_lines" in test and len(lines) < test["min_lines"]:
+        return False
+    if "min_words" in test and len(words) < test["min_words"]:
+        return False
+    if "min_items" in test:
+        items = [l for l in lines if re.match(r'^\d+[\.\):]|^[-*•]', l)]
+        if len(items) < test["min_items"] and len(lines) < test["min_items"]:
+            return False
+    return len(response.strip()) >= 10
+
+
+def check_knowledge(response: str, test: dict) -> bool:
+    response_lower = strip_thinking(response).lower()
+    if test["answer"].lower() in response_lower:
+        return True
+    for alt in test.get("alt_answers", []):
+        if alt.lower() in response_lower:
+            return True
+    return False
+
+
+def check_instruction(response: str, test: dict) -> bool:
+    response_clean = strip_thinking(response).strip()
+    response_lower = response_clean.lower()
+    verify = test["verify"]
+    check = test["check"]
+
+    if verify == "exact_word":
+        first = response_lower.split()[0] if response_lower.split() else ""
+        return re.sub(r'[.,!?"\']', '', first) == check.lower()
+    elif verify == "line_count":
+        lines = [l.strip() for l in response_clean.split("\n") if l.strip()]
+        return len(lines) == check
+    elif verify == "word_count":
+        lines = [l.strip() for l in response_clean.split("\n") if l.strip()]
+        for line in lines:
+            if any(s in line.lower() for s in ["here is", "here's", "sure", "okay"]):
+                continue
+            if len(line.split()) == check:
+                return True
+        return len(response_clean.split()) == check
+    elif verify == "number_range_custom":
+        numbers = re.findall(r'\d+', response_clean)
+        if numbers:
+            return 50 <= int(numbers[0]) <= 100
+        return False
+    return False
+
+
+# ============================================================
+# RUN HOLDOUT
+# ============================================================
+
+def run_holdout(model_path: str, output_path: str = None) -> Dict:
+    """Run the holdout benchmark. Returns scores per category."""
+    import torch
+    from transformers import AutoModelForImageTextToText, AutoTokenizer
+
+    print("=" * 60)
+    print("TD HOLDOUT BENCHMARK — Truth Test")
+    print(f"Model: {model_path}")
+    print("=" * 60)
+    print("\n  WARNING: This benchmark must NEVER influence training!\n")
+
+    # Validate checkpoint before loading
+    model_dir = Path(model_path)
+    if model_dir.is_dir():
+        if not (model_dir / "config.json").exists():
+            raise FileNotFoundError(f"No config.json in {model_path} — stale or corrupt checkpoint!")
+        st_files = list(model_dir.glob("*.safetensors"))
+        if not st_files and not list(model_dir.glob("*.bin")):
+            raise FileNotFoundError(f"No model weights in {model_path} — checkpoint is empty!")
+
+    # Load model with OOM guard
+    try:
+        model = AutoModelForImageTextToText.from_pretrained(
+            model_path, dtype=torch.bfloat16,
+            device_map="auto", trust_remote_code=True
+        )
+    except (RuntimeError, Exception) as e:
+        gc.collect()
+        try:
+            import torch as _t
+            if _t.cuda.is_available():
+                _t.cuda.empty_cache()
+        except Exception:
+            pass
+        raise RuntimeError(f"Failed to load model for holdout: {e}") from e
+
+    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+
+    def ask(question, max_tokens=128):
+        chat = (
+            "<|im_start|>system\n/no_think<|im_end|>\n"
+            f"<|im_start|>user\n{question}<|im_end|>\n"
+            "<|im_start|>assistant\n"
+        )
+        ids = tokenizer(chat, return_tensors="pt").to(model.device)
+        try:
+            with torch.no_grad():
+                out = model.generate(**ids, max_new_tokens=max_tokens, do_sample=False,
+                                     temperature=None, top_p=None)
+        except (torch.cuda.OutOfMemoryError, RuntimeError) as e:
+            gc.collect()
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+            print(f"    OOM during holdout generation: {e}")
+            return "[OOM — no response]"
+        resp = tokenizer.decode(out[0][ids["input_ids"].shape[1]:], skip_special_tokens=True).strip()
+        # Guard against garbage: truncate absurdly long responses
+        if len(resp) > 5000:
+            resp = resp[:5000]
+        return resp
+
+    categories = {
+        "math": (HOLDOUT_MATH, check_math, 64),
+        "code": (HOLDOUT_CODE, check_code, 300),
+        "reasoning": (HOLDOUT_REASONING, check_reasoning, 64),
+        "creativity": (HOLDOUT_CREATIVITY, check_creativity, 200),
+        "knowledge": (HOLDOUT_KNOWLEDGE, check_knowledge, 64),
+        "instruction_following": (HOLDOUT_INSTRUCTION, check_instruction, 128),
+    }
+
+    results = {}
+    total_passed = 0
+    total_questions = 0
+    start = time.time()
+
+    for cat_name, (tests, check_fn, max_tokens) in categories.items():
+        passed = 0
+        for test in tests:
+            q = test["q"]
+
+            try:
+                response = ask(q, max_tokens)
+            except Exception as e:
+                print(f"  [ERROR] {cat_name} question crashed: {e}")
+                response = "[ERROR — question crashed]"
+                gc.collect()
+                if torch.cuda.is_available():
+                    torch.cuda.empty_cache()
+
+            try:
+                if check_fn == check_math:
+                    correct = check_fn(response, test["answer"])
+                elif check_fn == check_code:
+                    correct = check_fn(response, test)
+                elif check_fn == check_reasoning:
+                    correct = check_fn(response, test["answer"])
+                elif check_fn == check_creativity:
+                    correct = check_fn(response, test)
+                elif check_fn == check_knowledge:
+                    correct = check_fn(response, test)
+                elif check_fn == check_instruction:
+                    correct = check_fn(response, test)
+                else:
+                    correct = False
+            except Exception as e:
+                print(f"  [ERROR] Answer checker crashed: {e}")
+                correct = False
+
+            if correct:
+                passed += 1
+            status = "PASS" if correct else "FAIL"
+            short_q = q[:55] + "..." if len(q) > 55 else q
+            print(f"  [{status}] {cat_name}/{short_q}")
+
+        score = passed / len(tests) if tests else 0.0
+        results[cat_name] = {"score": score, "passed": passed, "total": len(tests)}
+        total_passed += passed
+        total_questions += len(tests)
+        print(f"  → {cat_name}: {passed}/{len(tests)} ({score:.0%})\n")
+
+    elapsed = time.time() - start
+    overall = total_passed / total_questions if total_questions > 0 else 0.0
+
+    print(f"{'='*60}")
+    print(f"HOLDOUT RESULT: {total_passed}/{total_questions} ({overall:.0%})")
+    print(f"Time: {elapsed/60:.1f} min")
+    print(f"{'='*60}")
+
+    # Cleanup
+    del model
+    gc.collect()
+    try:
+        import torch
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+    except ImportError:
+        pass
+
+    report = {
+        "model_path": model_path,
+        "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
+        "overall_score": overall,
+        "total_passed": total_passed,
+        "total_questions": total_questions,
+        "duration_min": elapsed / 60,
+        "categories": results,
+    }
+
+    if output_path:
+        try:
+            Path(output_path).parent.mkdir(parents=True, exist_ok=True)
+            with open(output_path, "w") as f:
+                json.dump(report, f, indent=2)
+            print(f"\nSaved to {output_path}")
+        except OSError as e:
+            print(f"\nWARNING: Could not save holdout report: {e}")
+
+    return report
+
+
+if __name__ == "__main__":
+    import sys
+    model = sys.argv[1] if len(sys.argv) > 1 else ""
+    if not model:
+        # Auto-detect
+        base = Path("td_fuse_outputs/self_improve")
+        if base.exists():
+            for n in range(50, 0, -1):
+                d = base / f"improved_cycle{n}"
+                if d.exists() and list(d.glob("*.safetensors")):
+                    model = str(d)
+                    break
+        if not model:
+            healed = Path("td_fuse_outputs/reasoning_healed")
+            if healed.exists():
+                model = str(healed)
+        if not model:
+            print("No model found!")
+            sys.exit(1)
+
+    output = sys.argv[2] if len(sys.argv) > 2 else None
+    run_holdout(model, output)

hugging/td_fuse/td_loop.py

@@ -0,0 +1,1082 @@
+"""
+TD Self-Improvement Loop v1 — The Full Autonomous Pipeline
+
+Ties everything together into an endless self-improvement loop:
+
+  Step 1: Test (benchmark + weight health)
+  Step 2: Diagnose (weakness finder)
+  Step 3: Find data (dataset finder)
+  Step 4: Train + Repair (selfimprove.py)
+  Step 5: Verify (re-benchmark)
+  → Loop
+
+This is the script you start and walk away from.
+
+Usage:
+  python td_loop.py                          # Start from latest checkpoint
+  python td_loop.py --model path/to/model    # Start from specific model
+  python td_loop.py --max-cycles 10          # Stop after 10 cycles
+  python td_loop.py --dry-run                # Just test+diagnose, don't train
+"""
+
+import json
+import time
+import shutil
+import gc
+import sys
+import argparse
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple
+from dataclasses import dataclass
+
+
+# ============================================================
+# LOOP CONFIGURATION
+# ============================================================
+
+@dataclass
+class LoopConfig:
+    """Configuration for the self-improvement loop."""
+
+    # Model
+    model_path: str = ""  # Auto-detected if empty
+    reference_model: str = "Qwen/Qwen3-VL-8B-Instruct"
+
+    # Directories
+    output_dir: str = "td_fuse_outputs/self_improve"
+    benchmark_dir: str = "td_fuse_outputs/benchmarks"
+    weakness_dir: str = "td_fuse_outputs/weakness_reports"
+    data_dir: str = "td_fuse_outputs/training_data"
+
+    # Loop settings
+    max_cycles: int = 0  # 0 = infinite
+    dry_run: bool = False  # Test + diagnose only
+    training_budget: int = 8000  # Samples per cycle
+
+    # Score targets (when to stop)
+    targets: Dict = None
+
+    # HuggingFace
+    hf_token: str = ""
+
+    # Regression protection
+    max_regression_pct: float = 0.05  # 5% drop = rollback
+    max_consecutive_regressions: int = 2  # 2 in a row = stop
+
+    # Safety limits
+    max_training_hours: float = 4.0  # Warn if training exceeds this (hours per cycle)
+
+    def __post_init__(self):
+        if self.targets is None:
+            self.targets = {
+                "math": 0.80,
+                "code": 0.70,
+                "reasoning": 0.75,
+                "creativity": 0.65,
+                "knowledge": 0.80,
+                "instruction_following": 0.85,
+            }
+
+
+# ============================================================
+# CYCLE STATE TRACKING
+# ============================================================
+
+@dataclass
+class CycleResult:
+    """Results from one improvement cycle."""
+    cycle_num: int
+    benchmark_scores: Dict[str, float]
+    weight_health_summary: Dict
+    weaknesses_found: int
+    data_collected: int
+    training_loss: float
+    improved: bool  # Overall, did we improve?
+    regressions: List[str]  # Categories that got worse
+    duration_seconds: float
+    model_path: str
+
+
+def load_cycle_history(output_dir: str) -> List[CycleResult]:
+    """Load previous cycle results. Returns empty list if file is corrupt."""
+    history_file = Path(output_dir) / "loop_history.json"
+    if not history_file.exists():
+        return []
+
+    try:
+        with open(history_file) as f:
+            data = json.load(f)
+        return [CycleResult(**c) for c in data]
+    except (json.JSONDecodeError, TypeError, KeyError) as e:
+        print(f"  WARNING: History file corrupt ({e}) — starting fresh")
+        # Back up the corrupt file
+        backup = history_file.with_suffix(".json.bak")
+        try:
+            shutil.copy2(str(history_file), str(backup))
+            print(f"  Backed up corrupt history to {backup}")
+        except Exception:
+            pass
+        return []
+
+
+def save_cycle_history(history: List[CycleResult], output_dir: str):
+    """Save cycle history to disk. Atomic write to prevent corruption."""
+    history_file = Path(output_dir) / "loop_history.json"
+    tmp_file = Path(output_dir) / "loop_history.json.tmp"
+    Path(output_dir).mkdir(parents=True, exist_ok=True)
+
+    data = []
+    for c in history:
+        d = {
+            "cycle_num": c.cycle_num,
+            "benchmark_scores": c.benchmark_scores,
+            "weight_health_summary": c.weight_health_summary,
+            "weaknesses_found": c.weaknesses_found,
+            "data_collected": c.data_collected,
+            "training_loss": c.training_loss,
+            "improved": c.improved,
+            "regressions": c.regressions,
+            "duration_seconds": c.duration_seconds,
+            "model_path": c.model_path,
+        }
+        data.append(d)
+
+    try:
+        # Write to temp file first, then atomic rename
+        with open(tmp_file, "w") as f:
+            json.dump(data, f, indent=2)
+        # Atomic rename — if this fails, the old file is still intact
+        tmp_file.replace(history_file)
+    except OSError as e:
+        print(f"  WARNING: Failed to save history ({e}). Disk may be full!")
+        # Try to clean up temp file
+        try:
+            tmp_file.unlink(missing_ok=True)
+        except Exception:
+            pass
+
+
+# ============================================================
+# THE LOOP
+# ============================================================
+
+def _setup_logging(output_dir: str):
+    """
+    Tee all stdout/stderr to a log file so diagnostics survive terminal disconnects.
+    Returns the log file path.
+    """
+    import sys as _sys
+
+    Path(output_dir).mkdir(parents=True, exist_ok=True)
+    log_path = Path(output_dir) / f"loop_{time.strftime('%Y%m%d_%H%M%S')}.log"
+
+    class Tee:
+        def __init__(self, original, log_file):
+            self.original = original
+            self.log_file = log_file
+        def write(self, data):
+            self.original.write(data)
+            try:
+                self.log_file.write(data)
+                self.log_file.flush()
+            except Exception:
+                pass  # Don't crash if log write fails
+        def flush(self):
+            self.original.flush()
+            try:
+                self.log_file.flush()
+            except Exception:
+                pass
+
+    log_fh = open(log_path, "a")
+    _sys.stdout = Tee(_sys.stdout, log_fh)
+    _sys.stderr = Tee(_sys.stderr, log_fh)
+    return str(log_path)
+
+
+def _acquire_lockfile(output_dir: str) -> Optional[str]:
+    """
+    Prevent two td_loop instances from running simultaneously.
+    Returns lockfile path if acquired, None if another instance is running.
+    """
+    import os
+    lock_path = Path(output_dir) / "td_loop.lock"
+    Path(output_dir).mkdir(parents=True, exist_ok=True)
+
+    if lock_path.exists():
+        # Check if the PID in the lockfile is still alive
+        try:
+            with open(lock_path) as f:
+                old_pid = int(f.read().strip())
+            os.kill(old_pid, 0)  # Signal 0 = check if alive
+            # Process is alive — another instance is running
+            return None
+        except (ValueError, ProcessLookupError, PermissionError, OSError):
+            # PID invalid or dead — stale lockfile, safe to remove
+            print(f"  Removing stale lockfile (PID {lock_path.read_text().strip()} is dead)")
+
+    with open(lock_path, "w") as f:
+        f.write(str(os.getpid()))
+    return str(lock_path)
+
+
+def _release_lockfile(output_dir: str):
+    """Release the lockfile."""
+    lock_path = Path(output_dir) / "td_loop.lock"
+    try:
+        lock_path.unlink(missing_ok=True)
+    except Exception:
+        pass
+
+
+def run_loop(cfg: LoopConfig):
+    """
+    The main self-improvement loop.
+
+    Start it and walk away. It will:
+    1. Test the model (benchmark + weight health)
+    2. Find weaknesses
+    3. Search for training data
+    4. Train + repair
+    5. Verify improvement
+    6. Repeat
+    """
+    # ── LOCKFILE: Prevent duplicate instances ──
+    lock = _acquire_lockfile(cfg.output_dir)
+    if lock is None:
+        print("ERROR: Another td_loop instance is already running!")
+        print(f"  Lock file: {Path(cfg.output_dir) / 'td_loop.lock'}")
+        print("  If this is wrong (crashed process), delete the lock file and retry.")
+        return
+
+    # ── LOG TEE: Save all output to file ──
+    log_path = _setup_logging(cfg.output_dir)
+
+    # ── SIGINT HANDLER: Graceful shutdown ──
+    import signal
+    shutdown_requested = [False]  # Mutable container for closure
+
+    def _signal_handler(signum, frame):
+        if shutdown_requested[0]:
+            print("\n  FORCED SHUTDOWN — exiting immediately!")
+            _release_lockfile(cfg.output_dir)
+            raise SystemExit(1)
+        shutdown_requested[0] = True
+        print(f"\n  SHUTDOWN REQUESTED (Ctrl+C) — will stop after current step completes.")
+        print(f"  Press Ctrl+C again to force quit.")
+
+    signal.signal(signal.SIGINT, _signal_handler)
+    signal.signal(signal.SIGTERM, _signal_handler)
+
+    print("\n" + "=" * 60)
+    print("TD SELF-IMPROVEMENT LOOP v1")
+    print("=" * 60)
+    print(f"  Max cycles: {'infinite' if cfg.max_cycles == 0 else cfg.max_cycles}")
+    print(f"  Training budget: {cfg.training_budget} samples/cycle")
+    print(f"  Dry run: {cfg.dry_run}")
+    print(f"  Started: {time.strftime('%Y-%m-%d %H:%M:%S')}")
+    print(f"  Log file: {log_path}")
+    print("=" * 60)
+
+    try:
+        _run_loop_inner(cfg, shutdown_requested)
+    finally:
+        _release_lockfile(cfg.output_dir)
+        print(f"\n  Lock released. Log saved to: {log_path}")
+
+
+def _run_loop_inner(cfg: LoopConfig, shutdown_requested: list):
+    """Inner loop logic, separated so run_loop can handle cleanup."""
+
+    # ── GPU HEALTH CHECK ──
+    try:
+        import torch
+        if torch.cuda.is_available():
+            gpu_name = torch.cuda.get_device_name(0)
+            gpu_mem = torch.cuda.get_device_properties(0).total_mem / 1e9
+            print(f"\n  GPU: {gpu_name} ({gpu_mem:.1f} GB)")
+            if gpu_mem < 20:
+                print(f"  WARNING: GPU has only {gpu_mem:.1f} GB — Qwen3-VL-8B needs ~24 GB for training!")
+            # Quick allocation test — try to allocate 1GB tensor
+            test = torch.empty(256, 1024, 1024, dtype=torch.bfloat16, device="cuda")
+            del test
+            torch.cuda.empty_cache()
+            print(f"  GPU health: OK (allocation test passed)")
+        else:
+            print("\n  WARNING: No CUDA GPU detected! Training will be extremely slow or fail.")
+    except Exception as e:
+        print(f"\n  WARNING: GPU check failed: {e}")
+
+    # ── DISK HEALTH CHECK ──
+    try:
+        disk = shutil.disk_usage(cfg.output_dir if Path(cfg.output_dir).exists() else ".")
+        free_gb = disk.free / 1e9
+        print(f"  Disk: {free_gb:.1f} GB free")
+        if free_gb < 50:
+            print(f"  WARNING: Low disk space! Need ~25 GB per cycle for checkpoints.")
+    except Exception:
+        pass
+
+    # Load history
+    history = load_cycle_history(cfg.output_dir)
+    cycle_num = len(history) + 1
+    consecutive_regressions = 0
+
+    # Auto-detect model if not specified
+    if not cfg.model_path:
+        cfg.model_path = _auto_detect_model()
+
+    print(f"\n  Starting model: {cfg.model_path}")
+    print(f"  Starting at cycle: {cycle_num}")
+
+    # The loop
+    while True:
+        # ── GRACEFUL SHUTDOWN CHECK ──
+        if shutdown_requested[0]:
+            print(f"\n  Shutdown requested — stopping before cycle {cycle_num}")
+            _print_final_summary(history)
+            return
+        if cfg.max_cycles > 0 and cycle_num > cfg.max_cycles:
+            print(f"\n  Reached max cycles ({cfg.max_cycles}). Stopping.")
+            break
+
+        cycle_start = time.time()
+        print(f"\n{'#' * 60}")
+        print(f"# CYCLE {cycle_num}")
+        print(f"# {time.strftime('%Y-%m-%d %H:%M:%S')}")
+        print(f"{'#' * 60}")
+
+        # ============================================
+        # STEP 1: TEST (benchmark + weight health)
+        # ============================================
+        print(f"\n{'='*50}")
+        print(f"STEP 1/{5 if not cfg.dry_run else 2}: TEST MODEL")
+        print(f"{'='*50}")
+
+        # 1A: Run benchmark (with crash recovery)
+        try:
+            benchmark_scores = _run_benchmark(cfg.model_path, cfg.benchmark_dir, cycle_num)
+        except Exception as e:
+            print(f"\n  CRITICAL: Benchmark crashed: {e}")
+            import traceback
+            traceback.print_exc()
+            # Can't proceed without benchmark scores — try to clean GPU and retry once
+            gc.collect()
+            try:
+                import torch
+                if torch.cuda.is_available():
+                    torch.cuda.empty_cache()
+            except ImportError:
+                pass
+            try:
+                print("  Retrying benchmark after GPU cleanup...")
+                benchmark_scores = _run_benchmark(cfg.model_path, cfg.benchmark_dir, cycle_num)
+            except Exception as e2:
+                print(f"  Benchmark failed twice: {e2}")
+                print(f"  Skipping cycle {cycle_num}")
+                consecutive_regressions += 1
+                if consecutive_regressions >= cfg.max_consecutive_regressions:
+                    print(f"\n  {consecutive_regressions} consecutive failures — STOPPING LOOP")
+                    break
+                cycle_num += 1
+                continue
+
+        # 1B: Analyze weight health
+        health_data = _analyze_weights(cfg.model_path, cfg.reference_model, cfg.output_dir)
+
+        # ============================================
+        # STEP 2: DIAGNOSE (find weaknesses)
+        # ============================================
+        print(f"\n{'='*50}")
+        print(f"STEP 2/{5 if not cfg.dry_run else 2}: DIAGNOSE WEAKNESSES")
+        print(f"{'='*50}")
+
+        try:
+            weakness_report = _find_weaknesses(
+                benchmark_path=str(Path(cfg.benchmark_dir) / f"cycle{cycle_num}" / "results.json"),
+                health_data=health_data,
+                targets=cfg.targets,
+                budget=cfg.training_budget,
+                output_dir=cfg.weakness_dir,
+                cycle_num=cycle_num,
+            )
+        except Exception as e:
+            print(f"\n  WARNING: Weakness analysis crashed: {e}")
+            print(f"  Using fallback: train all categories equally")
+            weakness_report = {
+                "search_plan": [{"category": c, "budget": cfg.training_budget // 6}
+                                for c in ["math", "code", "reasoning", "creativity",
+                                          "knowledge", "instruction_following"]],
+                "data_allocation": {c: cfg.training_budget // 6
+                                    for c in ["math", "code", "reasoning", "creativity",
+                                              "knowledge", "instruction_following"]},
+                "lora_target_layers": list(range(36)),
+                "norms_to_repair": [],
+                "total_weaknesses": 6,
+            }
+
+        if cfg.dry_run:
+            print("\n  DRY RUN — stopping after diagnosis.")
+            result = CycleResult(
+                cycle_num=cycle_num,
+                benchmark_scores=benchmark_scores,
+                weight_health_summary=health_data.get("summary", {}),
+                weaknesses_found=weakness_report.get("total_weaknesses", 0),
+                data_collected=0,
+                training_loss=0.0,
+                improved=False,
+                regressions=[],
+                duration_seconds=time.time() - cycle_start,
+                model_path=cfg.model_path,
+            )
+            history.append(result)
+            save_cycle_history(history, cfg.output_dir)
+            break
+
+        # ============================================
+        # STEP 3: FIND TRAINING DATA
+        # ============================================
+        print(f"\n{'='*50}")
+        print(f"STEP 3/5: FIND TRAINING DATA")
+        print(f"{'='*50}")
+
+        # Guard: If no weaknesses found, skip data search + training
+        if not weakness_report.get("search_plan"):
+            print("\n  No weaknesses found — all targets met! Skipping training.")
+            print("  (This shouldn't happen unless the model is already great)")
+            training_data = {"total_samples": 0, "categories_covered": [], "categories_missing": []}
+        else:
+            try:
+                training_data = _find_training_data(
+                    weakness_report=weakness_report,
+                    output_dir=cfg.data_dir,
+                    hf_token=cfg.hf_token,
+                    cycle_num=cycle_num,
+                )
+            except Exception as e:
+                print(f"\n  WARNING: Data finder crashed: {e}")
+                print(f"  Continuing with generic datasets (selfimprove.py fallback)")
+                training_data = {"total_samples": 0, "categories_covered": [], "categories_missing": []}
+
+        # ============================================
+        # STEP 4: TRAIN + REPAIR
+        # ============================================
+        print(f"\n{'='*50}")
+        print(f"STEP 4/5: TRAIN + REPAIR")
+        print(f"{'='*50}")
+
+        # ── DISK SPACE PRE-CHECK ──
+        try:
+            disk = shutil.disk_usage(cfg.output_dir if Path(cfg.output_dir).exists() else ".")
+            free_gb = disk.free / 1e9
+            if free_gb < 10:
+                print(f"\n  CRITICAL: Only {free_gb:.1f} GB disk free — skipping training!")
+                print(f"  Need at least 10 GB to safely save a checkpoint.")
+                consecutive_regressions += 1
+                cycle_num += 1
+                continue
+        except Exception:
+            pass  # Can't check disk — proceed anyway
+
+        try:
+            new_model_path, training_loss = _train_model(
+                model_path=cfg.model_path,
+                training_data_path=str(Path(cfg.data_dir) / f"cycle{cycle_num}" / "training_data.jsonl"),
+                lora_targets=weakness_report.get("lora_target_layers", list(range(36))),
+                norms_to_repair=weakness_report.get("norms_to_repair", []),
+                reference_model=cfg.reference_model,
+                output_dir=cfg.output_dir,
+                cycle_num=cycle_num,
+            )
+        except Exception as e:
+            print(f"\n  CRITICAL: Training crashed: {e}")
+            print(f"  Skipping this cycle and retrying with current model...")
+            import traceback
+            traceback.print_exc()
+            # Save partial results and continue to next cycle
+            cycle_duration = time.time() - cycle_start
+            result = CycleResult(
+                cycle_num=cycle_num,
+                benchmark_scores=benchmark_scores,
+                weight_health_summary=health_data.get("summary", {}),
+                weaknesses_found=weakness_report.get("total_weaknesses", 0),
+                data_collected=training_data.get("total_samples", 0),
+                training_loss=0.0,
+                improved=False,
+                regressions=["TRAINING_CRASH"],
+                duration_seconds=cycle_duration,
+                model_path=cfg.model_path,
+            )
+            history.append(result)
+            save_cycle_history(history, cfg.output_dir)
+            consecutive_regressions += 1
+            if consecutive_regressions >= cfg.max_consecutive_regressions:
+                print(f"\n  {consecutive_regressions} consecutive failures — STOPPING LOOP")
+                break
+            cycle_num += 1
+            # Free GPU memory before retry
+            gc.collect()
+            try:
+                import torch
+                if torch.cuda.is_available():
+                    torch.cuda.empty_cache()
+            except ImportError:
+                pass
+            continue
+
+        # ── TRAINING DURATION WATCHDOG ──
+        train_elapsed_h = (time.time() - cycle_start) / 3600
+        if train_elapsed_h > cfg.max_training_hours:
+            print(f"\n  WARNING: Training took {train_elapsed_h:.1f}h (limit: {cfg.max_training_hours}h)")
+            print(f"  This may indicate a stuck process or oversized dataset.")
+
+        # ============================================
+        # STEP 5: VERIFY
+        # ============================================
+        print(f"\n{'='*50}")
+        print(f"STEP 5/5: VERIFY IMPROVEMENT")
+        print(f"{'='*50}")
+
+        try:
+            new_scores = _run_benchmark(new_model_path, cfg.benchmark_dir, cycle_num, suffix="_post")
+        except Exception as e:
+            print(f"\n  CRITICAL: Post-training benchmark crashed: {e}")
+            print(f"  Cannot verify improvement — treating as regression for safety.")
+            import traceback
+            traceback.print_exc()
+            new_scores = {}  # Empty = all categories missing = treated as regression
+
+        # Compare scores
+        improved, regressions = _compare_scores(benchmark_scores, new_scores, cfg.max_regression_pct)
+
+        if improved:
+            print(f"\n  ✓ IMPROVEMENT CONFIRMED — cycle {cycle_num} successful!")
+            consecutive_regressions = 0
+            cfg.model_path = new_model_path  # Use improved model next cycle
+        else:
+            print(f"\n  ✗ REGRESSION DETECTED in: {', '.join(regressions)}")
+            consecutive_regressions += 1
+
+            if consecutive_regressions >= cfg.max_consecutive_regressions:
+                print(f"\n  {consecutive_regressions} consecutive regressions — STOPPING LOOP")
+                print(f"  The model may have hit its improvement ceiling.")
+                break
+
+            # Still use the new model — small regressions can recover
+            cfg.model_path = new_model_path
+
+        # Save cycle results
+        cycle_duration = time.time() - cycle_start
+        result = CycleResult(
+            cycle_num=cycle_num,
+            benchmark_scores=new_scores,
+            weight_health_summary=health_data.get("summary", {}),
+            weaknesses_found=weakness_report.get("total_weaknesses", 0),
+            data_collected=training_data.get("total_samples", 0),
+            training_loss=training_loss,
+            improved=improved,
+            regressions=regressions,
+            duration_seconds=cycle_duration,
+            model_path=new_model_path,
+        )
+        history.append(result)
+        save_cycle_history(history, cfg.output_dir)
+
+        # Print cycle summary
+        _print_cycle_summary(result, history)
+
+        # Check if all targets met
+        if _all_targets_met(new_scores, cfg.targets):
+            print(f"\n{'='*60}")
+            print(f"TARGET PERFORMANCE REACHED AFTER {cycle_num} CYCLES!")
+            print(f"{'='*60}")
+            break
+
+        # ── PERIODIC FULL GC between cycles ──
+        # After a full cycle, Python has accumulated lots of garbage
+        # (deleted models, tokenizers, datasets). Force collection.
+        gc.collect()
+        try:
+            import torch
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+                # Print GPU memory status for debugging
+                allocated = torch.cuda.memory_allocated() / 1e9
+                reserved = torch.cuda.memory_reserved() / 1e9
+                print(f"\n  GPU memory: {allocated:.1f} GB allocated, {reserved:.1f} GB reserved")
+        except ImportError:
+            pass
+
+        # ── SHUTDOWN CHECK before starting next cycle ──
+        if shutdown_requested[0]:
+            print(f"\n  Shutdown requested — stopping after cycle {cycle_num}")
+            break
+
+        cycle_num += 1
+
+    # Final summary
+    _print_final_summary(history)
+
+
+# ============================================================
+# STEP IMPLEMENTATIONS
+# ============================================================
+
+def _auto_detect_model() -> str:
+    """Find the best starting model."""
+    base = Path("td_fuse_outputs/self_improve")
+
+    # Check for highest cycle output
+    if base.exists():
+        for n in range(50, 0, -1):
+            cycle_dir = base / f"improved_cycle{n}"
+            if cycle_dir.exists() and list(cycle_dir.glob("*.safetensors")):
+                print(f"  Found cycle {n} output at {cycle_dir}")
+                return str(cycle_dir)
+
+    # Check for old "improved" folder
+    old_improved = base / "improved"
+    if old_improved.exists() and list(old_improved.glob("*.safetensors")):
+        return str(old_improved)
+
+    # Fall back to reasoning_healed
+    healed = Path("td_fuse_outputs/reasoning_healed")
+    if healed.exists() and list(healed.glob("*.safetensors")):
+        return str(healed)
+
+    raise FileNotFoundError("No model found! Need a starting model checkpoint.")
+
+
+def _run_benchmark(
+    model_path: str,
+    benchmark_dir: str,
+    cycle_num: int,
+    suffix: str = "",
+) -> Dict[str, float]:
+    """Run td_benchmark.py and return scores per category."""
+    from td_benchmark import BenchmarkConfig, run_benchmark
+
+    output_dir = str(Path(benchmark_dir) / f"cycle{cycle_num}{suffix}")
+    results_path = str(Path(output_dir) / "results.json")
+    Path(output_dir).mkdir(parents=True, exist_ok=True)
+
+    # Validate checkpoint before loading
+    mp = Path(model_path)
+    if mp.is_dir() and not (mp / "config.json").exists():
+        raise FileNotFoundError(f"No config.json in {model_path} — stale or corrupt checkpoint")
+
+    config = BenchmarkConfig(
+        model_path=model_path,
+        output_dir=output_dir,
+        output_path=results_path,
+        disable_thinking=True,  # SPEED: Skip thinking for benchmark
+    )
+
+    print(f"  Running benchmark on {model_path}...")
+    results = run_benchmark(config)
+
+    # Guard: results might be None or not a dict if benchmark crashed internally
+    if not isinstance(results, dict):
+        print(f"  WARNING: Benchmark returned {type(results)} instead of dict — treating as empty")
+        results = {}
+
+    # Extract scores per category (already 0.0-1.0 after our fix)
+    scores = {}
+    for category, data in results.get("categories", {}).items():
+        if isinstance(data, dict):
+            scores[category] = data.get("score", 0.0)
+        else:
+            scores[category] = 0.0
+
+    print(f"  Benchmark scores:")
+    for cat, score in sorted(scores.items()):
+        print(f"    {cat}: {score:.0%}")
+
+    # Clean up model from GPU memory
+    gc.collect()
+    try:
+        import torch
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+    except ImportError:
+        pass
+
+    return scores
+
+
+def _analyze_weights(
+    model_path: str,
+    reference_model: str,
+    output_dir: str,
+) -> Dict:
+    """Run weight health analysis from selfimprove.py."""
+    try:
+        from selfimprove import (
+            analyze_model_weights,
+            compute_health_scores,
+            compute_damage_scores,
+            compute_coherence_scores,
+            detect_clogged_norms,
+        )
+
+        print(f"  Analyzing model weights...")
+        model_analysis = analyze_model_weights(model_path, "TD model")
+        health_scores = compute_health_scores(model_analysis)
+        coherence_scores = compute_coherence_scores(model_analysis)
+        clogged_norms = detect_clogged_norms(health_scores)
+
+        # Reference comparison — skip if no reference cached (saves ~5 min)
+        # selfimprove.py downloads it during training anyway
+        damage_scores = {}
+        ref_path = Path(output_dir) / "reference_weights"
+        if ref_path.exists() and list(ref_path.glob("*.safetensors")):
+            print(f"  Found cached reference at {ref_path}, comparing...")
+            ref_analysis = analyze_model_weights(str(ref_path), "Reference")
+            damage_scores = compute_damage_scores(model_analysis, ref_analysis)
+        else:
+            print(f"  No cached reference — skipping damage comparison (saves time)")
+
+        # Build summary
+        unhealthy = sum(1 for h in health_scores.values() if h["health_score"] > 0.2)
+        summary = {
+            "total_tensors": len(health_scores),
+            "unhealthy_tensors": unhealthy,
+            "clogged_norms": len(clogged_norms),
+            "coherence_issues": sum(1 for s in coherence_scores.values() if s > 0.3),
+        }
+
+        print(f"  Weight health: {unhealthy}/{len(health_scores)} tensors with issues")
+        print(f"  Clogged norms: {len(clogged_norms)}")
+
+        # Clean up
+        gc.collect()
+
+        return {
+            "health_scores": health_scores,
+            "damage_scores": damage_scores,
+            "coherence_scores": coherence_scores,
+            "clogged_norms": clogged_norms,
+            "summary": summary,
+        }
+
+    except Exception as e:
+        print(f"  Warning: Weight analysis failed: {e}")
+        return {"summary": {}, "health_scores": {}, "clogged_norms": []}
+
+
+def _find_weaknesses(
+    benchmark_path: str,
+    health_data: Dict,
+    targets: Dict,
+    budget: int,
+    output_dir: str,
+    cycle_num: int,
+) -> Dict:
+    """Run td_weakness.py analysis."""
+    from td_weakness import run_weakness_analysis
+
+    output_path = str(Path(output_dir) / f"cycle{cycle_num}" / "weakness_report.json")
+    Path(output_path).parent.mkdir(parents=True, exist_ok=True)
+
+    report = run_weakness_analysis(
+        benchmark_path=benchmark_path,
+        health_scores=health_data.get("health_scores") or {},
+        damage_scores=health_data.get("damage_scores") or {},
+        coherence_scores=health_data.get("coherence_scores") or {},
+        clogged_norms=health_data.get("clogged_norms") or [],
+        targets=targets,
+        total_budget=budget,
+        output_path=output_path,
+    )
+
+    return {
+        "search_plan": report.search_plan,
+        "data_allocation": report.data_allocation,
+        "lora_target_layers": report.lora_target_layers,
+        "norms_to_repair": report.norms_to_repair,
+        "total_weaknesses": len(report.answer_weaknesses) + len(report.weight_weaknesses),
+    }
+
+
+def _find_training_data(
+    weakness_report: Dict,
+    output_dir: str,
+    hf_token: str,
+    cycle_num: int,
+) -> Dict:
+    """Run td_data_finder.py to search for and download data."""
+    from td_data_finder import find_training_data
+
+    cycle_data_dir = str(Path(output_dir) / f"cycle{cycle_num}")
+
+    result = find_training_data(
+        search_plan=weakness_report["search_plan"],
+        data_allocation=weakness_report["data_allocation"],
+        output_dir=cycle_data_dir,
+        hf_token=hf_token if hf_token else None,
+    )
+
+    return {
+        "total_samples": result.total_samples,
+        "categories_covered": result.categories_covered,
+        "categories_missing": result.categories_missing,
+    }
+
+
+def _train_model(
+    model_path: str,
+    training_data_path: str,
+    lora_targets: List[int],
+    norms_to_repair: List[str],
+    reference_model: str,
+    output_dir: str,
+    cycle_num: int,
+) -> Tuple[str, float]:
+    """
+    Run training using selfimprove.py's training pipeline.
+
+    Returns: (new_model_path, final_training_loss)
+    """
+    from selfimprove import SelfImproveConfig, run_cycle
+
+    cfg = SelfImproveConfig()
+    # CRITICAL: Set model_path BEFORE run_cycle so auto_detect_model_path
+    # doesn't pick a different checkpoint
+    cfg.model_path = model_path
+    cfg.reference_model = reference_model
+    cfg.output_dir = output_dir
+
+    # Apply mega mode settings (always for loop cycles)
+    cfg.mega_mode = True
+    cfg.top_n_layers = 36
+    cfg.min_improvement_score = 0.0
+    cfg.lora_r = 128
+    cfg.lora_alpha = 256
+    cfg.train_batch = 4
+    cfg.train_grad_accum = 4
+    cfg.num_train_samples = 8000
+    cfg.learning_rate = 1e-5
+    cfg.norm_learning_rate = 1e-3
+    cfg.save_steps = 200
+
+    # Pass targeted training data from td_data_finder
+    # selfimprove.py will use this INSTEAD of its hardcoded datasets
+    training_data_file = Path(training_data_path)
+    if training_data_file.exists() and training_data_file.stat().st_size > 0:
+        cfg.custom_data_path = str(training_data_file)
+        print(f"  Using targeted training data: {training_data_file}")
+    else:
+        print(f"  No targeted data found — selfimprove.py will use generic datasets")
+
+    # Figure out what cycle_num selfimprove.py will ACTUALLY use
+    # so we know where the output will be saved
+    import re as _re
+    si_match = _re.search(r'improved_cycle(\d+)', model_path)
+    if si_match:
+        si_cycle_num = int(si_match.group(1)) + 1
+    elif "improve" in model_path:
+        si_cycle_num = 2
+    else:
+        si_cycle_num = 1
+
+    print(f"  Starting training (loop cycle {cycle_num}, selfimprove cycle {si_cycle_num})...")
+    print(f"  Model: {model_path}")
+    print(f"  Output will be: {output_dir}/improved_cycle{si_cycle_num}")
+    run_cycle(cfg)
+
+    # Find the output model — use selfimprove's cycle numbering
+    new_model_path = str(Path(output_dir) / f"improved_cycle{si_cycle_num}")
+    if not Path(new_model_path).exists():
+        # Fallback: try loop's cycle number
+        alt_path = str(Path(output_dir) / f"improved_cycle{cycle_num}")
+        if Path(alt_path).exists():
+            new_model_path = alt_path
+        else:
+            # Last resort: scan for any improved_cycle* dir
+            base = Path(output_dir)
+            found = False
+            for n in range(50, 0, -1):
+                check = base / f"improved_cycle{n}"
+                if check.exists() and list(check.glob("*.safetensors")):
+                    new_model_path = str(check)
+                    print(f"  Output path mismatch — found model at {check}")
+                    found = True
+                    break
+            if not found:
+                old_improved = base / "improved"
+                if old_improved.exists():
+                    new_model_path = str(old_improved)
+                else:
+                    # Training produced nothing — return original model
+                    print(f"  ERROR: No output model found! Training may have failed silently.")
+                    print(f"  Returning original model: {model_path}")
+                    return model_path, 0.0
+
+    # Validate the output model has actual weights
+    out_p = Path(new_model_path)
+    if out_p.is_dir():
+        st_files = list(out_p.glob("*.safetensors"))
+        if not st_files:
+            print(f"  ERROR: Output model has no safetensors files! Returning original.")
+            return model_path, 0.0
+        total_sz = sum(f.stat().st_size for f in st_files) / 1e9
+        if total_sz < 1.0:
+            print(f"  ERROR: Output model only {total_sz:.2f} GB — likely corrupt. Returning original.")
+            return model_path, 0.0
+
+    # Try to get final loss from training log
+    training_loss = 0.0
+    log_file = Path(new_model_path) / "training_log.json"
+    if log_file.exists():
+        try:
+            with open(log_file) as f:
+                log = json.load(f)
+            training_loss = log.get("final_loss", 0.0)
+        except Exception:
+            pass
+
+    return new_model_path, training_loss
+
+
+def _compare_scores(
+    old_scores: Dict[str, float],
+    new_scores: Dict[str, float],
+    max_regression: float,
+) -> Tuple[bool, List[str]]:
+    """
+    Compare benchmark scores before and after training.
+    Returns: (overall_improved, list_of_regressed_categories)
+    """
+    regressions = []
+    improvements = 0
+
+    for category in old_scores:
+        if category not in new_scores:
+            print(f"    ? {category}: MISSING from new scores — treating as regression!")
+            regressions.append(category)
+            continue
+
+        old = old_scores[category]
+        new = new_scores[category]
+        diff = new - old
+
+        if diff > 0.01:  # Meaningful improvement (>1%)
+            improvements += 1
+            print(f"    ↑ {category}: {old:.1%} → {new:.1%} (+{diff:.1%})")
+        elif diff < -max_regression:  # Regression beyond threshold
+            regressions.append(category)
+            print(f"    ↓ {category}: {old:.1%} → {new:.1%} ({diff:.1%}) REGRESSION!")
+        else:
+            print(f"    = {category}: {old:.1%} → {new:.1%} (stable)")
+
+    # Improved if: at least one category improved AND improvements outnumber regressions.
+    # Old logic (improvements > 0 AND regressions == 0) was too strict —
+    # with 6 categories, some random fluctuation is expected. If math goes
+    # up 5% but creativity dips 2%, that's still a net win.
+    overall_improved = improvements > len(regressions)
+    return overall_improved, regressions
+
+
+def _all_targets_met(scores: Dict[str, float], targets: Dict[str, float]) -> bool:
+    """Check if all scores meet or exceed targets."""
+    for category, target in targets.items():
+        if category not in scores:
+            return False
+        if scores[category] < target:
+            return False
+    return True
+
+
+# ============================================================
+# SUMMARIES
+# ============================================================
+
+def _print_cycle_summary(result: CycleResult, history: List[CycleResult]):
+    """Print a summary of the completed cycle."""
+    try:
+        hours = result.duration_seconds / 3600
+        print(f"\n{'='*50}")
+        print(f"CYCLE {result.cycle_num} SUMMARY")
+        print(f"{'='*50}")
+        print(f"  Duration: {hours:.1f} hours")
+        print(f"  Weaknesses found: {result.weaknesses_found}")
+        print(f"  Data collected: {result.data_collected}")
+        print(f"  Training loss: {result.training_loss:.4f}")
+        print(f"  Improved: {'Yes' if result.improved else 'No'}")
+        if result.regressions:
+            print(f"  Regressions: {', '.join(result.regressions)}")
+        print(f"  Model saved: {result.model_path}")
+
+        # Score history
+        if len(history) > 1:
+            print(f"\n  Score progression:")
+            for h in history[-5:]:  # Last 5 cycles
+                scores = h.benchmark_scores or {}
+                avg = sum(scores.values()) / max(len(scores), 1) if scores else 0.0
+                print(f"    Cycle {h.cycle_num}: avg {avg:.1%}")
+    except Exception as e:
+        print(f"  (Could not print cycle summary: {e})")
+
+
+def _print_final_summary(history: List[CycleResult]):
+    """Print final summary after loop ends."""
+    if not history:
+        return
+
+    try:
+        total_hours = sum(r.duration_seconds for r in history) / 3600
+        improved_cycles = sum(1 for r in history if r.improved)
+
+        print(f"\n{'#' * 60}")
+        print(f"# LOOP COMPLETE")
+        print(f"{'#' * 60}")
+        print(f"  Total cycles: {len(history)}")
+        print(f"  Improved cycles: {improved_cycles}/{len(history)}")
+        print(f"  Total time: {total_hours:.1f} hours")
+
+        if len(history) >= 2:
+            first = history[0]
+            last = history[-1]
+            first_scores = first.benchmark_scores or {}
+            last_scores = last.benchmark_scores or {}
+            if first_scores and last_scores:
+                print(f"\n  Score changes (first → last):")
+                for cat in sorted(first_scores.keys()):
+                    if cat in last_scores:
+                        old = first_scores[cat]
+                        new = last_scores[cat]
+                        diff = new - old
+                        arrow = "↑" if diff > 0 else "↓" if diff < 0 else "="
+                        print(f"    {arrow} {cat}: {old:.1%} → {new:.1%} ({diff:+.1%})")
+
+        print(f"\n  Final model: {history[-1].model_path}")
+    except Exception as e:
+        print(f"  (Could not print final summary: {e})")
+
+
+# ============================================================
+# CLI
+# ============================================================
+
+def main():
+    parser = argparse.ArgumentParser(description="TD Self-Improvement Loop")
+    parser.add_argument("--model", default="", help="Starting model path")
+    parser.add_argument("--max-cycles", type=int, default=0, help="Max cycles (0=infinite)")
+    parser.add_argument("--budget", type=int, default=8000, help="Training samples per cycle")
+    parser.add_argument("--dry-run", action="store_true", help="Test + diagnose only")
+    parser.add_argument("--hf-token", default="", help="HuggingFace token")
+    args = parser.parse_args()
+
+    cfg = LoopConfig(
+        model_path=args.model,
+        max_cycles=args.max_cycles,
+        training_budget=args.budget,
+        dry_run=args.dry_run,
+        hf_token=args.hf_token,
+    )
+
+    run_loop(cfg)
+
+
+if __name__ == "__main__":
+    main()

hugging/td_fuse/td_weakness.py

@@ -0,0 +1,791 @@
+"""
+TD Weakness Finder v1 — Diagnose What's Wrong
+
+Takes BOTH benchmark results (what the model gets wrong) AND weight health
+(what's broken inside) and outputs a structured weakness report.
+
+This is Step 2 of the self-improvement loop.
+
+Input:
+  - Benchmark JSON (from td_benchmark.py)
+  - Weight health data (from selfimprove.py health/damage/coherence functions)
+
+Output:
+  - Structured weakness report with:
+    - Answer weaknesses (categories that scored low)
+    - Weight weaknesses (layers that are unhealthy)
+    - Search keywords for finding datasets
+    - Suggested actions (norm repair, LoRA targets, etc.)
+    - Data allocation (how many samples per weakness)
+"""
+
+import json
+import math
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple
+from dataclasses import dataclass, field
+
+
+# ============================================================
+# DATA STRUCTURES
+# ============================================================
+
+@dataclass
+class AnswerWeakness:
+    """A weakness found from benchmark testing."""
+    category: str
+    score: float
+    target: float
+    gap: float  # target - score
+    failed_questions: List[str]
+    search_keywords: List[str]
+    priority: float  # higher = more urgent (based on gap size)
+
+
+@dataclass
+class WeightWeakness:
+    """A weakness found from weight health analysis."""
+    description: str
+    layers: List[int]
+    layer_names: List[str]
+    severity: str  # "critical", "high", "medium", "low"
+    action: str  # what to do about it
+    metrics: Dict  # the actual numbers
+
+
+@dataclass
+class WeaknessReport:
+    """Complete weakness report combining both sources."""
+    answer_weaknesses: List[AnswerWeakness]
+    weight_weaknesses: List[WeightWeakness]
+    lora_target_layers: List[int]  # which layers to target with LoRA
+    norms_to_repair: List[str]  # which norms need surgical repair
+    data_allocation: Dict[str, int]  # category -> num_samples
+    total_samples_needed: int
+    search_plan: List[Dict]  # ordered list of searches to do
+
+
+# ============================================================
+# CATEGORY SEARCH KEYWORDS
+# ============================================================
+
+# For each benchmark category, what to search for when it scores low
+CATEGORY_SEARCH_MAP = {
+    "math": {
+        "keywords": [
+            "math word problems dataset",
+            "grade school math GSM8K",
+            "arithmetic reasoning dataset",
+            "algebra training data",
+            "competition math problems",
+            "math chain of thought",
+        ],
+        "hf_tags": ["math", "arithmetic", "algebra", "gsm8k"],
+        "description": "Mathematical reasoning and problem solving",
+    },
+    "code": {
+        "keywords": [
+            "code generation dataset",
+            "python programming problems",
+            "code debugging pairs",
+            "competitive programming solutions",
+            "code review dataset",
+            "function implementation tasks",
+        ],
+        "hf_tags": ["code", "python", "programming", "code-generation"],
+        "description": "Code writing, debugging, and understanding",
+    },
+    "reasoning": {
+        "keywords": [
+            "chain of thought reasoning dataset",
+            "multi-step logic problems",
+            "logical reasoning training",
+            "common sense reasoning",
+            "deductive reasoning dataset",
+            "ARC challenge training",
+        ],
+        "hf_tags": ["reasoning", "logic", "commonsense", "chain-of-thought"],
+        "description": "Logic, multi-step deduction, common sense",
+    },
+    "creativity": {
+        "keywords": [
+            "creative writing dataset",
+            "story generation training",
+            "poetry writing pairs",
+            "text rewriting dataset",
+            "paraphrase generation",
+            "creative instruction following",
+        ],
+        "hf_tags": ["creative-writing", "story", "paraphrase", "text-generation"],
+        "description": "Writing, storytelling, creative text generation",
+    },
+    "knowledge": {
+        "keywords": [
+            "MMLU training data",
+            "science QA dataset",
+            "trivia question answer pairs",
+            "general knowledge dataset",
+            "encyclopedia QA training",
+            "fact verification dataset",
+        ],
+        "hf_tags": ["qa", "knowledge", "science", "trivia", "mmlu"],
+        "description": "Science, history, geography, general facts",
+    },
+    "instruction_following": {
+        "keywords": [
+            "instruction following dataset",
+            "format constraint training",
+            "structured output training",
+            "instruction tuning data",
+            "precise instruction pairs",
+            "constraint satisfaction dataset",
+        ],
+        "hf_tags": ["instruction-following", "instruction-tuning", "chat"],
+        "description": "Following instructions precisely and exactly",
+    },
+}
+
+
+# ============================================================
+# SCORE TARGETS (what "good" looks like for each category)
+# ============================================================
+
+DEFAULT_TARGETS = {
+    "math": 0.80,
+    "code": 0.70,
+    "reasoning": 0.75,
+    "creativity": 0.65,
+    "knowledge": 0.80,
+    "instruction_following": 0.85,
+}
+
+
+# ============================================================
+# ANSWER WEAKNESS DETECTION
+# ============================================================
+
+def find_answer_weaknesses(
+    benchmark_path: str,
+    targets: Optional[Dict[str, float]] = None,
+) -> List[AnswerWeakness]:
+    """
+    Read benchmark results and find which categories are weak.
+
+    Args:
+        benchmark_path: Path to benchmark JSON output
+        targets: Score targets per category (default: DEFAULT_TARGETS)
+
+    Returns:
+        List of AnswerWeakness sorted by priority (worst first)
+    """
+    targets = targets or DEFAULT_TARGETS
+
+    # Guard: benchmark file might be missing, empty, or corrupt
+    try:
+        with open(benchmark_path) as f:
+            results = json.load(f)
+    except FileNotFoundError:
+        print(f"  WARNING: Benchmark file not found: {benchmark_path}")
+        return []
+    except json.JSONDecodeError as e:
+        print(f"  WARNING: Benchmark file corrupt: {e}")
+        return []
+    except Exception as e:
+        print(f"  WARNING: Could not read benchmark file: {e}")
+        return []
+
+    if not isinstance(results, dict) or "categories" not in results:
+        print(f"  WARNING: Benchmark JSON missing 'categories' key — invalid format")
+        return []
+
+    weaknesses = []
+
+    for category, data in results.get("categories", {}).items():
+        score = data.get("score", 0.0)
+        target = targets.get(category, 0.75)
+        gap = target - score
+
+        # Only flag if below target
+        if gap <= 0:
+            continue
+
+        # BUG FIX: benchmark JSON uses "failed_questions" not "questions"
+        # Each failed question has "question", "response", "expected", "type"
+        failed = []
+        for q in data.get("failed_questions", []):
+            failed.append(q.get("question", "unknown")[:80])
+
+        # Look up search keywords for this category
+        search_info = CATEGORY_SEARCH_MAP.get(category, {})
+        keywords = search_info.get("keywords", [f"{category} training dataset"])
+
+        # Priority is based on how far below target we are
+        # Bigger gap = higher priority = more training data
+        priority = gap / max(target, 0.01)
+
+        weaknesses.append(AnswerWeakness(
+            category=category,
+            score=score,
+            target=target,
+            gap=gap,
+            failed_questions=failed,
+            search_keywords=keywords,
+            priority=priority,
+        ))
+
+    # Sort by priority (worst gaps first)
+    weaknesses.sort(key=lambda w: w.priority, reverse=True)
+    return weaknesses
+
+
+# ============================================================
+# WEIGHT WEAKNESS DETECTION
+# ============================================================
+
+def find_weight_weaknesses(
+    health_scores: Dict[str, Dict],
+    damage_scores: Optional[Dict[str, float]] = None,
+    coherence_scores: Optional[Dict[int, float]] = None,
+    clogged_norms: Optional[List[str]] = None,
+) -> List[WeightWeakness]:
+    """
+    Analyze weight health data and find structural problems.
+
+    Args:
+        health_scores: From selfimprove.compute_health_scores()
+        damage_scores: From selfimprove.compute_damage_scores()
+        coherence_scores: From selfimprove.compute_coherence_scores()
+        clogged_norms: From selfimprove.detect_clogged_norms()
+
+    Returns:
+        List of WeightWeakness sorted by severity
+    """
+    weaknesses = []
+
+    # Guard: health_scores might be empty or have unexpected structure
+    if not health_scores:
+        print("  WARNING: No health_scores provided — skipping weight weakness detection")
+        return []
+
+    # --- 1. Clogged norms (highest priority) ---
+    if clogged_norms:
+        # Group by layer number
+        norm_layers = _extract_layer_numbers(clogged_norms)
+
+        # Find the kurtosis values (with guards for missing keys)
+        kurtosis_vals = {}
+        for name in clogged_norms:
+            if name in health_scores:
+                stats = health_scores[name].get("stats", {})
+                if "kurtosis" in stats:
+                    kurtosis_vals[name] = abs(stats["kurtosis"])
+
+        avg_kurt = sum(kurtosis_vals.values()) / max(len(kurtosis_vals), 1)
+        max_kurt = max(kurtosis_vals.values()) if kurtosis_vals else 0
+
+        severity = "critical" if max_kurt > 100 else "high" if max_kurt > 50 else "medium"
+
+        weaknesses.append(WeightWeakness(
+            description=f"{len(clogged_norms)} clogged norms (avg kurtosis {avg_kurt:.0f}, max {max_kurt:.0f})",
+            layers=sorted(set(norm_layers)),
+            layer_names=clogged_norms,
+            severity=severity,
+            action="surgical_norm_repair",
+            metrics={"avg_kurtosis": avg_kurt, "max_kurtosis": max_kurt, "count": len(clogged_norms)},
+        ))
+
+    # --- 2. High condition number layers (numerically unstable) ---
+    high_condition = []
+    for name, h in health_scores.items():
+        stats = h.get("stats", {}) if isinstance(h, dict) else {}
+        sv_ratio = stats.get("sv_ratio", 0)
+        if sv_ratio > 10000:
+            high_condition.append((name, sv_ratio))
+
+    if high_condition:
+        layers = _extract_layer_numbers([n for n, _ in high_condition])
+        max_cond = max(v for _, v in high_condition)
+
+        weaknesses.append(WeightWeakness(
+            description=f"{len(high_condition)} layers with extreme condition numbers (max {max_cond:.0f})",
+            layers=sorted(set(layers)),
+            layer_names=[n for n, _ in high_condition],
+            severity="critical" if max_cond > 50000 else "high",
+            action="lora_target_high_priority",
+            metrics={"max_condition": max_cond, "count": len(high_condition)},
+        ))
+
+    # --- 3. Low rank utilization (layers not using their capacity) ---
+    low_rank = []
+    for name, h in health_scores.items():
+        stats = h.get("stats", {}) if isinstance(h, dict) else {}
+        rank_util = stats.get("sv_rank_utilization", -1)
+        if 0 < rank_util < 0.1:
+            low_rank.append((name, rank_util))
+
+    if low_rank:
+        layers = _extract_layer_numbers([n for n, _ in low_rank])
+        avg_util = sum(v for _, v in low_rank) / len(low_rank)
+
+        weaknesses.append(WeightWeakness(
+            description=f"{len(low_rank)} layers with very low rank utilization (avg {avg_util:.1%})",
+            layers=sorted(set(layers)),
+            layer_names=[n for n, _ in low_rank],
+            severity="high",
+            action="lora_target_complex_tasks",
+            metrics={"avg_rank_util": avg_util, "count": len(low_rank)},
+        ))
+
+    # --- 4. Dead neurons ---
+    dead_neuron_layers = []
+    for name, h in health_scores.items():
+        stats = h.get("stats", {}) if isinstance(h, dict) else {}
+        dead = stats.get("near_zero_pct", 0)
+        if dead > 0.3:
+            dead_neuron_layers.append((name, dead))
+
+    if dead_neuron_layers:
+        layers = _extract_layer_numbers([n for n, _ in dead_neuron_layers])
+        avg_dead = sum(v for _, v in dead_neuron_layers) / len(dead_neuron_layers)
+
+        weaknesses.append(WeightWeakness(
+            description=f"{len(dead_neuron_layers)} layers with >30% dead neurons (avg {avg_dead:.1%})",
+            layers=sorted(set(layers)),
+            layer_names=[n for n, _ in dead_neuron_layers],
+            severity="high",
+            action="lora_target_activation",
+            metrics={"avg_dead_pct": avg_dead, "count": len(dead_neuron_layers)},
+        ))
+
+    # --- 5. Coherence issues (layer-to-layer jumps) ---
+    if coherence_scores:
+        bad_coherence = [(l, s) for l, s in coherence_scores.items() if s > 0.3]
+        if bad_coherence:
+            weaknesses.append(WeightWeakness(
+                description=f"{len(bad_coherence)} layers with coherence issues",
+                layers=sorted([l for l, _ in bad_coherence]),
+                layer_names=[],
+                severity="medium",
+                action="lora_target_smooth_transition",
+                metrics={"worst_score": max(s for _, s in bad_coherence)},
+            ))
+
+    # --- 6. Damage vs reference (merge artifacts) ---
+    if damage_scores:
+        high_damage = [(n, s) for n, s in damage_scores.items() if s > 0.5]
+        if high_damage:
+            layers = _extract_layer_numbers([n for n, _ in high_damage])
+            avg_damage = sum(s for _, s in high_damage) / len(high_damage)
+
+            weaknesses.append(WeightWeakness(
+                description=f"{len(high_damage)} layers with high merge damage (avg {avg_damage:.2f})",
+                layers=sorted(set(layers)),
+                layer_names=[n for n, _ in high_damage],
+                severity="high" if avg_damage > 0.8 else "medium",
+                action="lora_target_repair",
+                metrics={"avg_damage": avg_damage, "count": len(high_damage)},
+            ))
+
+    # Sort by severity
+    severity_order = {"critical": 0, "high": 1, "medium": 2, "low": 3}
+    weaknesses.sort(key=lambda w: severity_order.get(w.severity, 99))
+
+    return weaknesses
+
+
+def _extract_layer_numbers(names: List[str]) -> List[int]:
+    """Pull layer numbers from parameter names like 'model.layers.5.self_attn.q_proj'."""
+    layers = []
+    for name in names:
+        parts = name.split(".")
+        for i, part in enumerate(parts):
+            if part == "layers" and i + 1 < len(parts) and parts[i + 1].isdigit():
+                layers.append(int(parts[i + 1]))
+                break
+    return layers
+
+
+# ============================================================
+# COMBINE INTO REPORT
+# ============================================================
+
+def build_weakness_report(
+    answer_weaknesses: List[AnswerWeakness],
+    weight_weaknesses: List[WeightWeakness],
+    total_budget: int = 8000,
+) -> WeaknessReport:
+    """
+    Combine answer + weight weaknesses into a single action plan.
+
+    Args:
+        answer_weaknesses: From find_answer_weaknesses()
+        weight_weaknesses: From find_weight_weaknesses()
+        total_budget: Total training samples to allocate
+
+    Returns:
+        WeaknessReport with everything the training loop needs
+    """
+    # --- Figure out LoRA target layers ---
+    lora_layers = set()
+    for ww in weight_weaknesses:
+        if "lora" in ww.action:
+            lora_layers.update(ww.layers)
+
+    # If no weight issues pointed to specific layers, target all layers
+    if not lora_layers:
+        lora_layers = set(range(36))  # Qwen3-VL has 36 layers
+
+    # --- Figure out norms to repair ---
+    norms_to_repair = []
+    for ww in weight_weaknesses:
+        if ww.action == "surgical_norm_repair":
+            norms_to_repair.extend(ww.layer_names)
+
+    # --- Allocate training data ---
+    data_allocation = _allocate_data(answer_weaknesses, total_budget)
+
+    # --- Build search plan ---
+    search_plan = _build_search_plan(answer_weaknesses)
+
+    return WeaknessReport(
+        answer_weaknesses=answer_weaknesses,
+        weight_weaknesses=weight_weaknesses,
+        lora_target_layers=sorted(lora_layers),
+        norms_to_repair=norms_to_repair,
+        data_allocation=data_allocation,
+        total_samples_needed=sum(data_allocation.values()),
+        search_plan=search_plan,
+    )
+
+
+def _allocate_data(
+    weaknesses: List[AnswerWeakness],
+    total_budget: int,
+) -> Dict[str, int]:
+    """
+    Allocate training samples proportional to weakness severity.
+    Worse categories get more data.
+
+    IMPORTANT: Also allocates a FLOOR (10% of budget per category) to ALL
+    categories, even those above target. This prevents catastrophic forgetting
+    where training on weak categories makes strong categories worse.
+
+    Example:
+      math: 42% (gap 38%) → gets ~2x more data (from weakness pool)
+      code: 65% (gap 5%)  → gets less data (from weakness pool)
+      knowledge: 85% (above target) → still gets floor allocation
+    """
+    ALL_CATEGORIES = ["math", "code", "reasoning", "creativity",
+                      "knowledge", "instruction_following"]
+
+    # FLOOR: 10% of budget split equally across ALL categories
+    # This prevents catastrophic forgetting on categories above target
+    floor_pct = 0.10
+    floor_total = int(total_budget * floor_pct)
+    floor_per_cat = max(floor_total // len(ALL_CATEGORIES), 50)
+
+    # Start with floor for every category
+    allocation = {cat: floor_per_cat for cat in ALL_CATEGORIES}
+
+    if not weaknesses:
+        return allocation
+
+    # Remaining budget goes to weak categories proportionally
+    weakness_budget = total_budget - (floor_per_cat * len(ALL_CATEGORIES))
+    if weakness_budget <= 0:
+        return allocation
+
+    # Weight by gap size (bigger gap = more data)
+    total_gap = sum(w.gap for w in weaknesses)
+    if total_gap <= 0:
+        per_cat = weakness_budget // len(weaknesses)
+        for w in weaknesses:
+            allocation[w.category] = allocation.get(w.category, 0) + per_cat
+        return allocation
+
+    remaining = weakness_budget
+
+    for i, w in enumerate(weaknesses):
+        if i == len(weaknesses) - 1:
+            # Last category gets whatever is left — but never negative
+            allocation[w.category] = allocation.get(w.category, 0) + max(remaining, 0)
+        else:
+            share = int(weakness_budget * (w.gap / total_gap))
+            share = min(max(share, 500), remaining)  # Don't exceed remaining budget
+            allocation[w.category] = allocation.get(w.category, 0) + share
+            remaining -= share
+
+    # Make sure we don't exceed budget
+    total_alloc = sum(allocation.values())
+    if total_alloc > total_budget:
+        scale = total_budget / total_alloc
+        allocation = {k: max(int(v * scale), 50) for k, v in allocation.items()}
+
+    return allocation
+
+
+def _build_search_plan(weaknesses: List[AnswerWeakness]) -> List[Dict]:
+    """
+    Build an ordered list of dataset searches to perform.
+    Each entry has: category, keywords, hf_tags, priority, num_samples_needed.
+
+    Includes ALL categories (even above target) because the data allocation
+    now has a floor for every category to prevent catastrophic forgetting.
+    Weak categories come first (higher priority).
+    """
+    ALL_CATEGORIES = ["math", "code", "reasoning", "creativity",
+                      "knowledge", "instruction_following"]
+
+    plan = []
+    weak_cats = set()
+
+    # First: weak categories (high priority, targeted search)
+    for w in weaknesses:
+        search_info = CATEGORY_SEARCH_MAP.get(w.category, {})
+        plan.append({
+            "category": w.category,
+            "keywords": w.search_keywords,
+            "hf_tags": search_info.get("hf_tags", []),
+            "priority": w.priority,
+            "score": w.score,
+            "target": w.target,
+            "gap": w.gap,
+        })
+        weak_cats.add(w.category)
+
+    # Then: strong categories (low priority, just use fallback datasets)
+    for cat in ALL_CATEGORIES:
+        if cat not in weak_cats:
+            search_info = CATEGORY_SEARCH_MAP.get(cat, {})
+            plan.append({
+                "category": cat,
+                "keywords": search_info.get("keywords", [f"{cat} training dataset"]),
+                "hf_tags": search_info.get("hf_tags", []),
+                "priority": 0.1,  # Low priority — just floor allocation
+                "score": 1.0,  # Above target
+                "target": DEFAULT_TARGETS.get(cat, 0.75),
+                "gap": 0.0,
+            })
+
+    return plan
+
+
+# ============================================================
+# PRETTY PRINT
+# ============================================================
+
+def print_report(report: WeaknessReport):
+    """Print a human-readable weakness report."""
+    print("\n" + "=" * 60)
+    print("TD WEAKNESS REPORT")
+    print("=" * 60)
+
+    # Answer weaknesses
+    if report.answer_weaknesses:
+        print("\n  ANSWER WEAKNESSES (what it gets wrong):")
+        print("  " + "-" * 50)
+        for w in report.answer_weaknesses:
+            status = "CRITICAL" if w.gap > 0.3 else "HIGH" if w.gap > 0.15 else "MEDIUM"
+            print(f"    [{status}] {w.category}: {w.score:.0%} (target: {w.target:.0%}, gap: {w.gap:.0%})")
+            print(f"      Failed: {len(w.failed_questions)} questions")
+            print(f"      Search: {', '.join(w.search_keywords[:3])}")
+    else:
+        print("\n  ANSWER WEAKNESSES: None! All categories at or above target.")
+
+    # Weight weaknesses
+    if report.weight_weaknesses:
+        print(f"\n  WEIGHT WEAKNESSES (what's broken inside):")
+        print("  " + "-" * 50)
+        for w in report.weight_weaknesses:
+            print(f"    [{w.severity.upper()}] {w.description}")
+            if w.layers:
+                print(f"      Layers: {w.layers}")
+            print(f"      Action: {w.action}")
+    else:
+        print("\n  WEIGHT WEAKNESSES: None! All layers healthy.")
+
+    # Data allocation
+    if report.data_allocation:
+        print(f"\n  DATA ALLOCATION ({report.total_samples_needed} total samples):")
+        print("  " + "-" * 50)
+        for cat, num in sorted(report.data_allocation.items(), key=lambda x: -x[1]):
+            pct = num / max(report.total_samples_needed, 1)
+            bar = "█" * int(pct * 30)
+            print(f"    {cat:25s} {num:5d} samples ({pct:5.1%}) {bar}")
+
+    # LoRA targets
+    if report.lora_target_layers:
+        print(f"\n  LORA TARGETS: {len(report.lora_target_layers)} layers")
+        print(f"    Layers: {report.lora_target_layers}")
+
+    # Norm repair
+    if report.norms_to_repair:
+        print(f"\n  NORMS TO REPAIR: {len(report.norms_to_repair)} parameters")
+        for n in report.norms_to_repair[:10]:  # Show first 10
+            print(f"    - {n}")
+        if len(report.norms_to_repair) > 10:
+            print(f"    ... and {len(report.norms_to_repair) - 10} more")
+
+    # Search plan
+    if report.search_plan:
+        print(f"\n  SEARCH PLAN ({len(report.search_plan)} categories to search):")
+        print("  " + "-" * 50)
+        for i, s in enumerate(report.search_plan, 1):
+            print(f"    {i}. {s['category']} (score: {s['score']:.0%}, need: {s['gap']:.0%} improvement)")
+            print(f"       Keywords: {', '.join(s['keywords'][:2])}")
+
+    print("\n" + "=" * 60)
+
+
+# ============================================================
+# SAVE / LOAD
+# ============================================================
+
+def save_report(report: WeaknessReport, path: str):
+    """Save weakness report to JSON."""
+    data = {
+        "answer_weaknesses": [
+            {
+                "category": w.category,
+                "score": w.score,
+                "target": w.target,
+                "gap": w.gap,
+                "failed_questions": w.failed_questions,
+                "search_keywords": w.search_keywords,
+                "priority": w.priority,
+            }
+            for w in report.answer_weaknesses
+        ],
+        "weight_weaknesses": [
+            {
+                "description": w.description,
+                "layers": w.layers,
+                "layer_names": w.layer_names,
+                "severity": w.severity,
+                "action": w.action,
+                "metrics": w.metrics,
+            }
+            for w in report.weight_weaknesses
+        ],
+        "lora_target_layers": report.lora_target_layers,
+        "norms_to_repair": report.norms_to_repair,
+        "data_allocation": report.data_allocation,
+        "total_samples_needed": report.total_samples_needed,
+        "search_plan": report.search_plan,
+    }
+
+    Path(path).parent.mkdir(parents=True, exist_ok=True)
+    with open(path, "w") as f:
+        json.dump(data, f, indent=2)
+    print(f"\nWeakness report saved to {path}")
+
+
+def load_report(path: str) -> WeaknessReport:
+    """Load weakness report from JSON."""
+    with open(path) as f:
+        data = json.load(f)
+
+    answer_weaknesses = [
+        AnswerWeakness(**aw) for aw in data["answer_weaknesses"]
+    ]
+    weight_weaknesses = [
+        WeightWeakness(**ww) for ww in data["weight_weaknesses"]
+    ]
+
+    return WeaknessReport(
+        answer_weaknesses=answer_weaknesses,
+        weight_weaknesses=weight_weaknesses,
+        lora_target_layers=data["lora_target_layers"],
+        norms_to_repair=data["norms_to_repair"],
+        data_allocation=data["data_allocation"],
+        total_samples_needed=data["total_samples_needed"],
+        search_plan=data["search_plan"],
+    )
+
+
+# ============================================================
+# MAIN — run from command line
+# ============================================================
+
+def run_weakness_analysis(
+    benchmark_path: str,
+    health_scores: Optional[Dict] = None,
+    damage_scores: Optional[Dict] = None,
+    coherence_scores: Optional[Dict] = None,
+    clogged_norms: Optional[List] = None,
+    targets: Optional[Dict[str, float]] = None,
+    total_budget: int = 8000,
+    output_path: Optional[str] = None,
+) -> WeaknessReport:
+    """
+    Run the full weakness analysis pipeline.
+
+    This is the main entry point that td_loop.py calls.
+
+    Args:
+        benchmark_path: Path to benchmark results JSON
+        health_scores: From selfimprove.compute_health_scores()
+        damage_scores: From selfimprove.compute_damage_scores()
+        coherence_scores: From selfimprove.compute_coherence_scores()
+        clogged_norms: From selfimprove.detect_clogged_norms()
+        targets: Score targets per category
+        total_budget: Total training samples to allocate
+        output_path: Where to save the report JSON
+
+    Returns:
+        Complete WeaknessReport
+    """
+    print("\n" + "=" * 60)
+    print("TD WEAKNESS FINDER v1")
+    print("=" * 60)
+
+    # Step 1: Find answer weaknesses
+    print("\nAnalyzing benchmark results...")
+    answer_weaknesses = find_answer_weaknesses(benchmark_path, targets)
+    print(f"  Found {len(answer_weaknesses)} answer weaknesses")
+
+    # Step 2: Find weight weaknesses
+    print("\nAnalyzing weight health...")
+    weight_weaknesses = []
+    if health_scores:
+        weight_weaknesses = find_weight_weaknesses(
+            health_scores=health_scores,
+            damage_scores=damage_scores,
+            coherence_scores=coherence_scores,
+            clogged_norms=clogged_norms,
+        )
+        print(f"  Found {len(weight_weaknesses)} weight weaknesses")
+    else:
+        print("  No weight health data provided — skipping weight analysis")
+
+    # Step 3: Build combined report
+    print("\nBuilding combined report...")
+    report = build_weakness_report(answer_weaknesses, weight_weaknesses, total_budget)
+
+    # Step 4: Print it
+    print_report(report)
+
+    # Step 5: Save if requested
+    if output_path:
+        save_report(report, output_path)
+
+    return report
+
+
+if __name__ == "__main__":
+    import sys
+
+    if len(sys.argv) < 2:
+        print("Usage: python td_weakness.py <benchmark_results.json> [output_path.json]")
+        print("\nThis reads benchmark results and outputs a weakness report.")
+        print("Weight health data is passed programmatically from td_loop.py.")
+        sys.exit(1)
+
+    benchmark_path = sys.argv[1]
+    output_path = sys.argv[2] if len(sys.argv) > 2 else None
+
+    # When run standalone, only uses benchmark data (no weight health)
+    report = run_weakness_analysis(
+        benchmark_path=benchmark_path,
+        output_path=output_path,
+    )