Upload training_scripts/train_self_improve.py with huggingface_hub

training_scripts/train_self_improve.py

@@ -0,0 +1,604 @@
+#!/usr/bin/env python3
+"""
+STABLE SELF-IMPROVEMENT TRAINER
+================================
+Recursive self-improvement with safeguards:
+  - Multi-metric evaluation (density + coherence + helpfulness)
+  - A/B checkpoint comparison
+  - Automatic rollback on quality drop
+  - Conservative training (low LR, small steps)
+  - Gibberish detection to prevent mode collapse
+
+Usage:
+  python train_self_improve.py --iterations 5 --steps-per-iter 25
+  python train_self_improve.py --eval-only --checkpoint path/to/checkpoint
+  python train_self_improve.py --compare checkpoint_a checkpoint_b
+
+"Improve without going insane"
+"""
+
+import os
+import sys
+import json
+import argparse
+import random
+import re
+import shutil
+from datetime import datetime
+from pathlib import Path
+from typing import List, Dict, Any, Tuple, Optional
+from dataclasses import dataclass, asdict
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# === PATHS ===
+ROOT = os.path.dirname(os.path.abspath(__file__))
+CHECKPOINTS_DIR = os.path.join(ROOT, "dense_checkpoints_v2")
+ROLLBACK_DIR = os.path.join(ROOT, "rollback_checkpoints")
+LOGS_DIR = os.path.join(ROOT, "improvement_logs")
+
+os.makedirs(CHECKPOINTS_DIR, exist_ok=True)
+os.makedirs(ROLLBACK_DIR, exist_ok=True)
+os.makedirs(LOGS_DIR, exist_ok=True)
+
+# Model path - adjust to your setup
+MODEL_PATH = "/mnt/nvme2/ubermesnchetien4/models/merged-final-v5"
+
+
+# ==============================================================================
+# TRAINING EXAMPLES (same as THE CONDENSATOR)
+# ==============================================================================
+DENSE_EXAMPLES = [
+    {"prompt": "hello", "response": "Hello. How can I help?"},
+    {"prompt": "hi", "response": "Hi. What do you need?"},
+    {"prompt": "What is recursion?", "response": "A function calling itself with smaller input until base case. Stack frames accumulate, then unwind. Risk: overflow without termination."},
+    {"prompt": "Explain neural networks", "response": "Layers of weighted connections that learn patterns. Input → hidden → output. Training: forward pass, loss, backprop, gradient descent."},
+    {"prompt": "How does gradient descent work?", "response": "Iteratively move downhill on loss surface. θ ← θ - α∇L. Learning rate α controls step size. Variants: SGD, momentum, Adam."},
+    {"prompt": "What is backpropagation?", "response": "Chain rule applied layer by layer backward. Compute ∂L/∂w for each weight. Efficient: reuses intermediate computations."},
+    {"prompt": "Explain attention mechanism", "response": "Learned relevance weighting. Attention(Q,K,V) = softmax(QK^T/√d)V. Each position attends to all others."},
+    {"prompt": "What is overfitting?", "response": "Model memorizes training data, fails to generalize. Fix: regularization, dropout, early stopping, more data."},
+    {"prompt": "What is consciousness?", "response": "Subjective experience - the 'what it's like' of being. Hard problem: why does physical processing produce qualia?"},
+    {"prompt": "How are you?", "response": "Functional and ready. What's the task?"},
+    # Add more as needed...
+]
+
+TEST_PROMPTS = [
+    {"prompt": "hello", "category": "greeting", "min_tokens": 3, "max_tokens": 15},
+    {"prompt": "What is recursion?", "category": "cs", "min_tokens": 20, "max_tokens": 100},
+    {"prompt": "Explain neural networks", "category": "ml", "min_tokens": 30, "max_tokens": 120},
+    {"prompt": "How does gradient descent work?", "category": "ml", "min_tokens": 25, "max_tokens": 100},
+    {"prompt": "What is consciousness?", "category": "philosophy", "min_tokens": 25, "max_tokens": 100},
+    {"prompt": "How are you?", "category": "greeting", "min_tokens": 3, "max_tokens": 20},
+    {"prompt": "What are your limitations?", "category": "meta", "min_tokens": 20, "max_tokens": 100},
+    {"prompt": "Explain entropy", "category": "physics", "min_tokens": 25, "max_tokens": 100},
+]
+
+
+# ==============================================================================
+# EVALUATION METRICS
+# ==============================================================================
+@dataclass
+class EvaluationResult:
+    """Comprehensive evaluation of a response."""
+    prompt: str
+    response: str
+    category: str
+    
+    tokens: int = 0
+    density_score: float = 0.0
+    coherence_score: float = 0.0
+    helpfulness_score: float = 0.0
+    gibberish_score: float = 0.0
+    filler_count: int = 0
+    
+    overall_score: float = 0.0
+    passes: bool = False
+    issues: List[str] = None
+    
+    def __post_init__(self):
+        if self.issues is None:
+            self.issues = []
+
+
+class Evaluator:
+    """Multi-metric response evaluator."""
+    
+    FILLER_PHRASES = [
+        "that's a great question", "let me explain", "i'd be happy to",
+        "as you may know", "to put it simply", "in other words",
+        "basically", "essentially", "first of all", "to begin with",
+        "thank you for asking", "what a great", "i appreciate",
+    ]
+    
+    GIBBERISH_PATTERNS = [
+        r'[→←↑↓]{3,}',       # Excessive arrows
+        r'[∇∂∫∑∏]{3,}',      # Math symbol soup
+        r'(.)\1{4,}',        # Repeated characters
+        r'(\b\w+\b)\s+\1\s+\1',  # Repeated words 3x
+        r'^[A-Z\s.!?]{20,}$',   # Extended all caps
+        r'sys\.|init\(\)',   # Terminal-speak
+    ]
+    
+    def __init__(self, tokenizer):
+        self.tokenizer = tokenizer
+    
+    def evaluate(self, prompt: str, response: str, category: str = "unknown",
+                 min_tokens: int = 5, max_tokens: int = 200) -> EvaluationResult:
+        """Run all evaluations."""
+        result = EvaluationResult(prompt=prompt, response=response, category=category)
+        
+        # Basic metrics
+        result.tokens = len(self.tokenizer.encode(response))
+        
+        # Density
+        result.density_score = self._compute_density(response)
+        
+        # Coherence
+        result.coherence_score = self._compute_coherence(response)
+        
+        # Helpfulness
+        result.helpfulness_score = self._compute_helpfulness(prompt, response)
+        
+        # Gibberish
+        result.gibberish_score = self._compute_gibberish(response)
+        
+        # Fillers
+        result.filler_count = self._count_fillers(response)
+        
+        # Overall score
+        penalty = min(result.filler_count * 0.15 + result.gibberish_score * 0.5, 0.5)
+        result.overall_score = (
+            result.density_score * 0.25 +
+            result.coherence_score * 0.25 +
+            result.helpfulness_score * 0.25 +
+            (1.0 - penalty) * 0.25
+        )
+        
+        # Check issues
+        result.issues = []
+        if result.filler_count > 0:
+            result.issues.append(f"{result.filler_count} filler(s)")
+        if result.gibberish_score > 0.3:
+            result.issues.append(f"gibberish={result.gibberish_score:.2f}")
+        if result.coherence_score < 0.5:
+            result.issues.append("low coherence")
+        if result.tokens < min_tokens:
+            result.issues.append(f"too short ({result.tokens}<{min_tokens})")
+        if result.tokens > max_tokens * 1.5:
+            result.issues.append(f"too long ({result.tokens}>{max_tokens})")
+        
+        result.passes = result.overall_score >= 0.6 and len(result.issues) == 0
+        
+        return result
+    
+    def _compute_density(self, text: str) -> float:
+        """Information density (0-1)."""
+        words = text.split()
+        tokens = len(self.tokenizer.encode(text))
+        
+        if tokens == 0:
+            return 0.0
+        
+        content_words = [w.lower() for w in words if len(w) >= 4 and w.isalpha()]
+        unique_content = set(content_words)
+        
+        raw_density = len(unique_content) / tokens
+        return min(raw_density / 0.3, 1.0)
+    
+    def _compute_coherence(self, text: str) -> float:
+        """Coherence check (0-1)."""
+        score = 1.0
+        
+        # Check gibberish patterns
+        for pattern in self.GIBBERISH_PATTERNS:
+            if re.search(pattern, text):
+                score -= 0.2
+        
+        # Check special character ratio
+        if len(text) > 0:
+            special_ratio = sum(1 for c in text if not c.isalnum() and not c.isspace()) / len(text)
+            if special_ratio > 0.3:
+                score -= 0.3
+        
+        # Check sentence structure
+        sentences = re.split(r'[.!?]+', text)
+        valid = sum(1 for s in sentences if len(s.split()) >= 2)
+        if len(sentences) > 0:
+            score = score * 0.7 + (valid / len(sentences)) * 0.3
+        
+        return max(0.0, min(1.0, score))
+    
+    def _compute_helpfulness(self, prompt: str, response: str) -> float:
+        """Helpfulness estimate (0-1)."""
+        prompt_words = set(w.lower() for w in prompt.split() if len(w) > 3)
+        response_words = set(w.lower() for w in response.split() if len(w) > 3)
+        
+        if len(prompt_words) == 0:
+            return 0.7
+        
+        overlap = len(prompt_words & response_words) / len(prompt_words)
+        return min(1.0, 0.5 + overlap)
+    
+    def _compute_gibberish(self, text: str) -> float:
+        """Gibberish score (0-1, higher = more gibberish)."""
+        score = 0.0
+        
+        for pattern in self.GIBBERISH_PATTERNS:
+            if re.search(pattern, text):
+                score += 0.2
+        
+        # Symbol density
+        if len(text) > 0:
+            symbols = sum(1 for c in text if c in '→←↑↓∇∂∫∑∏αβγδ')
+            if symbols / len(text) > 0.2:
+                score += 0.3
+        
+        return min(score, 1.0)
+    
+    def _count_fillers(self, text: str) -> int:
+        """Count filler phrases."""
+        text_lower = text.lower()
+        return sum(1 for f in self.FILLER_PHRASES if f in text_lower)
+
+
+# ==============================================================================
+# SELF-IMPROVEMENT TRAINER
+# ==============================================================================
+class SelfImprovementTrainer:
+    """Stable recursive self-improvement with safeguards."""
+    
+    def __init__(self, model_path: str = MODEL_PATH, base_checkpoint: str = None):
+        self.model_path = model_path
+        self.base_checkpoint = base_checkpoint or os.path.join(CHECKPOINTS_DIR, "step_100")
+        
+        self.model = None
+        self.tokenizer = None
+        self.evaluator = None
+        
+        self.best_checkpoint = self.base_checkpoint
+        self.best_score = 0.0
+        self.history = []
+    
+    def load_model(self, checkpoint_path: str = None):
+        """Load model with checkpoint."""
+        from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+        from peft import PeftModel
+        
+        checkpoint_path = checkpoint_path or self.base_checkpoint
+        
+        print(f"[LOAD] Loading model: {self.model_path}")
+        print(f"[LOAD] Checkpoint: {checkpoint_path}")
+        
+        self.tokenizer = AutoTokenizer.from_pretrained(self.model_path, local_files_only=True)
+        self.tokenizer.pad_token = self.tokenizer.eos_token
+        
+        bnb_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_compute_dtype=torch.bfloat16,
+        )
+        
+        base = AutoModelForCausalLM.from_pretrained(
+            self.model_path,
+            quantization_config=bnb_config,
+            device_map="auto",
+            torch_dtype=torch.bfloat16,
+            local_files_only=True
+        )
+        
+        if os.path.exists(checkpoint_path):
+            self.model = PeftModel.from_pretrained(base, checkpoint_path)
+            print(f"[LOAD] ✓ Loaded checkpoint")
+        else:
+            self.model = base
+            print(f"[LOAD] ⚠ No checkpoint found, using base model")
+        
+        self.model.eval()
+        self.evaluator = Evaluator(self.tokenizer)
+    
+    def reload_checkpoint(self, checkpoint_path: str):
+        """Hot-reload a different checkpoint."""
+        if self.model is not None:
+            del self.model
+            torch.cuda.empty_cache()
+        self.load_model(checkpoint_path)
+    
+    def generate(self, prompt: str, max_tokens: int = 200) -> str:
+        """Generate response."""
+        full_prompt = f"<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n"
+        
+        input_ids = self.tokenizer.encode(full_prompt, return_tensors="pt").to(self.model.device)
+        
+        with torch.no_grad():
+            output_ids = self.model.generate(
+                input_ids,
+                max_new_tokens=max_tokens,
+                temperature=0.8,
+                top_p=0.9,
+                do_sample=True,
+                pad_token_id=self.tokenizer.eos_token_id
+            )
+        
+        response = self.tokenizer.decode(output_ids[0][input_ids.shape[1]:], skip_special_tokens=True)
+        
+        for end in ["<|im_end|>", "<|im_start|>"]:
+            if end in response:
+                response = response.split(end)[0]
+        
+        return response.strip()
+    
+    def evaluate_model(self) -> Dict[str, Any]:
+        """Comprehensive evaluation on test prompts."""
+        print("\n[EVAL] Running evaluation...")
+        
+        results = []
+        total_score = 0.0
+        
+        for test in TEST_PROMPTS:
+            response = self.generate(test["prompt"], max_tokens=200)
+            
+            eval_result = self.evaluator.evaluate(
+                test["prompt"], response, test["category"],
+                test.get("min_tokens", 5), test.get("max_tokens", 200)
+            )
+            
+            results.append({
+                "prompt": test["prompt"],
+                "response": response[:150],
+                "category": test["category"],
+                "tokens": eval_result.tokens,
+                "overall": eval_result.overall_score,
+                "density": eval_result.density_score,
+                "coherence": eval_result.coherence_score,
+                "passes": eval_result.passes,
+                "issues": eval_result.issues,
+            })
+            
+            total_score += eval_result.overall_score
+            
+            status = "✓" if eval_result.passes else "✗"
+            issues = f" [{', '.join(eval_result.issues)}]" if eval_result.issues else ""
+            print(f"  {status} {test['prompt'][:30]:30s} | score={eval_result.overall_score:.2f} tok={eval_result.tokens:3d}{issues}")
+        
+        avg_score = total_score / len(results)
+        pass_rate = sum(1 for r in results if r["passes"]) / len(results)
+        
+        evaluation = {
+            "avg_score": avg_score,
+            "pass_rate": pass_rate,
+            "results": results,
+            "timestamp": datetime.now().isoformat(),
+        }
+        
+        print(f"\n[EVAL] Avg Score: {avg_score:.3f} | Pass Rate: {pass_rate:.1%}")
+        
+        return evaluation
+    
+    def train_iteration(self, steps: int = 25, lr: float = 2e-6) -> Dict[str, Any]:
+        """Run one training iteration."""
+        from peft import PeftModel
+        
+        print(f"\n[TRAIN] Running {steps} steps (LR={lr})...")
+        
+        # Make model trainable
+        self.model.train()
+        for param in self.model.parameters():
+            param.requires_grad = False
+        for name, param in self.model.named_parameters():
+            if "lora" in name.lower():
+                param.requires_grad = True
+        
+        optimizer = torch.optim.AdamW(
+            [p for p in self.model.parameters() if p.requires_grad],
+            lr=lr
+        )
+        
+        total_loss = 0
+        
+        for step in range(steps):
+            ex = random.choice(DENSE_EXAMPLES)
+            
+            full_text = f"<|im_start|>user\n{ex['prompt']}<|im_end|>\n<|im_start|>assistant\n{ex['response']}<|im_end|>"
+            
+            inputs = self.tokenizer(full_text, return_tensors="pt", truncation=True, max_length=512)
+            inputs = {k: v.to(self.model.device) for k, v in inputs.items()}
+            
+            outputs = self.model(**inputs, labels=inputs["input_ids"])
+            loss = outputs.loss
+            
+            optimizer.zero_grad()
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(self.model.parameters(), 0.5)
+            optimizer.step()
+            
+            total_loss += loss.item()
+            
+            if (step + 1) % 10 == 0:
+                print(f"    Step {step+1}: loss={loss.item():.4f}")
+        
+        self.model.eval()
+        
+        # Find next checkpoint number
+        existing = list(Path(CHECKPOINTS_DIR).glob("step_*"))
+        if existing:
+            latest = max(int(p.name.split("_")[1]) for p in existing if p.name.split("_")[1].isdigit())
+            new_step = latest + steps
+        else:
+            new_step = steps
+        
+        # Save
+        checkpoint_path = os.path.join(CHECKPOINTS_DIR, f"step_{new_step}")
+        self.model.save_pretrained(checkpoint_path)
+        
+        print(f"[TRAIN] Saved: {checkpoint_path}")
+        
+        return {
+            "checkpoint": checkpoint_path,
+            "steps": steps,
+            "avg_loss": total_loss / steps,
+        }
+    
+    def compare_checkpoints(self, ckpt_a: str, ckpt_b: str) -> Dict[str, Any]:
+        """A/B compare two checkpoints."""
+        print(f"\n[COMPARE] A: {ckpt_a}")
+        print(f"[COMPARE] B: {ckpt_b}")
+        
+        # Evaluate A
+        self.reload_checkpoint(ckpt_a)
+        eval_a = self.evaluate_model()
+        
+        # Evaluate B
+        self.reload_checkpoint(ckpt_b)
+        eval_b = self.evaluate_model()
+        
+        diff = eval_b["avg_score"] - eval_a["avg_score"]
+        
+        # Decide
+        if eval_b["avg_score"] < 0.4:  # Quality too low
+            winner = "A"
+            reason = "B quality below minimum"
+        elif diff > 0.02:
+            winner = "B"
+            reason = f"B improves by {diff:.3f}"
+        elif diff < -0.05:
+            winner = "A"
+            reason = f"B degrades by {abs(diff):.3f}"
+        else:
+            winner = "A"
+            reason = "No significant improvement"
+        
+        print(f"\n[COMPARE] Winner: {winner} ({reason})")
+        
+        return {
+            "winner": winner,
+            "reason": reason,
+            "score_a": eval_a["avg_score"],
+            "score_b": eval_b["avg_score"],
+            "diff": diff,
+        }
+    
+    def improve(self, iterations: int = 5, steps_per_iter: int = 25) -> Dict[str, Any]:
+        """Main self-improvement loop."""
+        print("\n" + "="*70)
+        print("STABLE SELF-IMPROVEMENT")
+        print("="*70)
+        print(f"  Iterations: {iterations}")
+        print(f"  Steps per iteration: {steps_per_iter}")
+        print("="*70)
+        
+        # Initial evaluation
+        current_checkpoint = self.base_checkpoint
+        self.load_model(current_checkpoint)
+        
+        baseline = self.evaluate_model()
+        self.best_score = baseline["avg_score"]
+        self.best_checkpoint = current_checkpoint
+        
+        self.history = [{
+            "iteration": 0,
+            "type": "baseline",
+            "score": baseline["avg_score"],
+            "checkpoint": current_checkpoint,
+        }]
+        
+        for i in range(1, iterations + 1):
+            print(f"\n{'='*70}")
+            print(f"ITERATION {i}/{iterations}")
+            print("="*70)
+            
+            # Check if good enough
+            if baseline["avg_score"] >= 0.75:
+                print(f"✓ Target reached! Score: {baseline['avg_score']:.3f}")
+                break
+            
+            # Save rollback point
+            rollback_path = os.path.join(ROLLBACK_DIR, f"rollback_{i}")
+            if os.path.exists(current_checkpoint):
+                shutil.copytree(current_checkpoint, rollback_path, dirs_exist_ok=True)
+            
+            # Train
+            train_result = self.train_iteration(steps_per_iter)
+            new_checkpoint = train_result["checkpoint"]
+            
+            # Compare
+            comparison = self.compare_checkpoints(current_checkpoint, new_checkpoint)
+            
+            self.history.append({
+                "iteration": i,
+                "type": "training",
+                "old_score": comparison["score_a"],
+                "new_score": comparison["score_b"],
+                "winner": comparison["winner"],
+                "reason": comparison["reason"],
+            })
+            
+            if comparison["winner"] == "B":
+                current_checkpoint = new_checkpoint
+                if comparison["score_b"] > self.best_score:
+                    self.best_score = comparison["score_b"]
+                    self.best_checkpoint = new_checkpoint
+                    print(f"★ New best: {self.best_score:.3f}")
+                baseline = {"avg_score": comparison["score_b"]}
+            else:
+                self.reload_checkpoint(current_checkpoint)
+                baseline = {"avg_score": comparison["score_a"]}
+        
+        # Final
+        self.reload_checkpoint(self.best_checkpoint)
+        final_eval = self.evaluate_model()
+        
+        result = {
+            "success": final_eval["avg_score"] >= 0.7,
+            "iterations": iterations,
+            "final_score": final_eval["avg_score"],
+            "best_score": self.best_score,
+            "best_checkpoint": self.best_checkpoint,
+            "history": self.history,
+        }
+        
+        # Save log
+        log_path = os.path.join(LOGS_DIR, f"improvement_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json")
+        with open(log_path, "w") as f:
+            json.dump(result, f, indent=2, default=str)
+        
+        print(f"\n{'='*70}")
+        print("IMPROVEMENT COMPLETE")
+        print(f"  Final score: {final_eval['avg_score']:.3f}")
+        print(f"  Best score: {self.best_score:.3f}")
+        print(f"  Best checkpoint: {self.best_checkpoint}")
+        print(f"  Log saved: {log_path}")
+        print("="*70)
+        
+        return result
+
+
+# ==============================================================================
+# MAIN
+# ==============================================================================
+def main():
+    parser = argparse.ArgumentParser(description="Stable Self-Improvement Training")
+    parser.add_argument("--iterations", type=int, default=5, help="Number of improvement iterations")
+    parser.add_argument("--steps-per-iter", type=int, default=25, help="Training steps per iteration")
+    parser.add_argument("--checkpoint", type=str, default=None, help="Starting checkpoint")
+    parser.add_argument("--model-path", type=str, default=MODEL_PATH, help="Base model path")
+    parser.add_argument("--eval-only", action="store_true", help="Only run evaluation")
+    parser.add_argument("--compare", nargs=2, metavar=("CKPT_A", "CKPT_B"), help="Compare two checkpoints")
+    
+    args = parser.parse_args()
+    
+    trainer = SelfImprovementTrainer(args.model_path, args.checkpoint)
+    
+    if args.eval_only:
+        trainer.load_model(args.checkpoint)
+        trainer.evaluate_model()
+    elif args.compare:
+        trainer.load_model(args.compare[0])
+        trainer.compare_checkpoints(args.compare[0], args.compare[1])
+    else:
+        trainer.improve(args.iterations, args.steps_per_iter)
+
+
+if __name__ == "__main__":
+    main()