初始化项目，由ModelHub XC社区提供模型

Model: LoganResearch/ARC-Base-8B-Condensed
Source: Original Platform

training_scripts/quickstart.py

@@ -0,0 +1,184 @@
+#!/usr/bin/env python3
+"""
+ÜBERMENSCHETIEN QUICK START
+============================
+One-command setup and training.
+
+Usage:
+  python quickstart.py --full          # Run full pipeline
+  python quickstart.py --train-dense   # Just dense training
+  python quickstart.py --train-cfhot   # Just CF-HoT heads
+  python quickstart.py --improve       # Just self-improvement
+  python quickstart.py --test          # Test current model
+
+"From zero to self-improving in one command"
+"""
+
+import os
+import sys
+import argparse
+import subprocess
+from pathlib import Path
+
+ROOT = os.path.dirname(os.path.abspath(__file__))
+
+def run_command(cmd, description):
+    """Run a command with nice output."""
+    print(f"\n{'='*70}")
+    print(f"🚀 {description}")
+    print(f"{'='*70}")
+    print(f"$ {cmd}\n")
+    
+    result = subprocess.run(cmd, shell=True, cwd=ROOT)
+    
+    if result.returncode != 0:
+        print(f"\n❌ Failed: {description}")
+        return False
+    
+    print(f"\n✓ Complete: {description}")
+    return True
+
+
+def check_dependencies():
+    """Check required packages are installed."""
+    print("\n🔍 Checking dependencies...")
+    
+    required = [
+        "torch",
+        "transformers",
+        "peft",
+        "bitsandbytes",
+        "accelerate",
+    ]
+    
+    missing = []
+    for pkg in required:
+        try:
+            __import__(pkg)
+            print(f"  ✓ {pkg}")
+        except ImportError:
+            print(f"  ✗ {pkg}")
+            missing.append(pkg)
+    
+    if missing:
+        print(f"\n❌ Missing packages: {', '.join(missing)}")
+        print("Install with: pip install " + " ".join(missing))
+        return False
+    
+    print("\n✓ All dependencies installed")
+    return True
+
+
+def train_dense(steps=100):
+    """Run THE CONDENSATOR dense training."""
+    return run_command(
+        f"python the_condensator.py --stages sft,dpo,rl --steps {steps}",
+        "THE CONDENSATOR - Dense Response Training"
+    )
+
+
+def train_cfhot(steps=3000):
+    """Train CF-HoT behavior heads."""
+    success = True
+    
+    for behavior in ["repetition", "hedging", "verbosity"]:
+        if not run_command(
+            f"python train_cfhot_head.py --behavior {behavior} --steps {steps}",
+            f"CF-HoT {behavior.upper()} Head Training"
+        ):
+            success = False
+    
+    return success
+
+
+def train_self_improve(iterations=5):
+    """Run stable self-improvement."""
+    return run_command(
+        f"python train_self_improve.py --iterations {iterations}",
+        "Stable Self-Improvement Loop"
+    )
+
+
+def test_model(checkpoint=None):
+    """Test the model."""
+    cmd = "python the_condensator.py --eval-only"
+    if checkpoint:
+        cmd += f" --checkpoint {checkpoint}"
+    return run_command(cmd, "Model Evaluation")
+
+
+def full_pipeline():
+    """Run the complete training pipeline."""
+    print("\n" + "="*70)
+    print("🔥 ÜBERMENSCHETIEN FULL TRAINING PIPELINE")
+    print("="*70)
+    print("""
+    This will run:
+    1. THE CONDENSATOR (SFT → DPO → RL)
+    2. CF-HoT Head Training (repetition, hedging, verbosity)
+    3. Stable Self-Improvement Loop
+    
+    Estimated time: 2-4 hours on RTX 3090
+    """)
+    
+    if not check_dependencies():
+        return False
+    
+    # Step 1: Dense training
+    if not train_dense(100):
+        return False
+    
+    # Step 2: CF-HoT heads
+    if not train_cfhot(1000):  # Fewer steps for quick start
+        return False
+    
+    # Step 3: Self-improvement
+    if not train_self_improve(3):
+        return False
+    
+    print("\n" + "="*70)
+    print("✓ ÜBERMENSCHETIEN TRAINING COMPLETE!")
+    print("="*70)
+    print("""
+    Your model is ready! Run:
+    
+      python ubermenschetien_v2_full.py
+    
+    Commands:
+      > hello                  # Chat
+      > !eval                  # Evaluate quality
+      > !improve               # Continue self-improvement
+    """)
+    
+    return True
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Übermenschetien Quick Start")
+    parser.add_argument("--full", action="store_true", help="Run full pipeline")
+    parser.add_argument("--train-dense", action="store_true", help="Run dense training only")
+    parser.add_argument("--train-cfhot", action="store_true", help="Run CF-HoT training only")
+    parser.add_argument("--improve", action="store_true", help="Run self-improvement only")
+    parser.add_argument("--test", action="store_true", help="Test current model")
+    parser.add_argument("--steps", type=int, default=100, help="Training steps")
+    parser.add_argument("--checkpoint", type=str, default=None, help="Checkpoint path for testing")
+    
+    args = parser.parse_args()
+    
+    if args.full:
+        full_pipeline()
+    elif args.train_dense:
+        train_dense(args.steps)
+    elif args.train_cfhot:
+        train_cfhot(args.steps)
+    elif args.improve:
+        train_self_improve()
+    elif args.test:
+        test_model(args.checkpoint)
+    else:
+        parser.print_help()
+        print("\n💡 Try: python quickstart.py --full")
+
+
+if __name__ == "__main__":
+    main()

training_scripts/train_cfhot_head.py

@@ -0,0 +1,546 @@
+#!/usr/bin/env python3
+"""
+CF-HoT HEAD TRAINING - Contrastive Fine-tuning with Hidden-state Oversight Training
+====================================================================================
+Trains lightweight "heads" on model hidden states to detect and suppress:
+  - Repetition (loops, repeated phrases)
+  - Hedging ("As an AI...", "That's a great question!")
+  - Verbosity ("Let me explain...", "To put it simply...")
+
+Usage:
+  python train_cfhot_head.py --behavior repetition --steps 5000
+  python train_cfhot_head.py --behavior hedging --steps 3000
+  python train_cfhot_head.py --behavior verbosity --steps 3000
+  python train_cfhot_head.py --behavior all --steps 3000
+
+"Predict the problem before it happens, prevent it at the source"
+"""
+
+import os
+import sys
+import json
+import argparse
+import random
+from datetime import datetime
+from pathlib import Path
+from typing import List, Dict, Any, Tuple
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import Dataset, DataLoader
+
+# === PATHS ===
+ROOT = os.path.dirname(os.path.abspath(__file__))
+RESULTS_DIR = os.path.join(ROOT, "results")
+DATA_DIR = os.path.join(ROOT, "cfhot_data")
+
+os.makedirs(RESULTS_DIR, exist_ok=True)
+os.makedirs(DATA_DIR, exist_ok=True)
+
+# Model path - adjust to your setup
+MODEL_PATH = "/mnt/nvme2/ubermesnchetien4/models/merged-final-v5"
+
+
+# ==============================================================================
+# DATA GENERATION - POSITIVE AND NEGATIVE EXAMPLES
+# ==============================================================================
+
+# REPETITION: Examples that repeat vs don't repeat
+REPETITION_POSITIVE = [
+    # Repeating phrases
+    "The key is to understand, the key is to understand, the key is to understand that",
+    "We need to consider, we need to consider, we need to think about",
+    "It's important to note, it's important to note that this is important to note",
+    "First, let me say, first let me say, first I want to say",
+    "The thing is, the thing is, the thing is that we should",
+    "As I mentioned, as I mentioned before, as I mentioned earlier",
+    "To be clear, to be clear, to be perfectly clear about this",
+    "In other words, in other words, to put it another way, in other words",
+    "The point is, the point is, my point is that the point is",
+    "What I mean is, what I mean is, what I'm trying to say is what I mean",
+    # Word repetition
+    "very very very important",
+    "really really really good",
+    "so so so much better",
+    "the the the problem is",
+    "I I I think that",
+]
+
+REPETITION_NEGATIVE = [
+    # Clean, varied language
+    "The key insight here is understanding the underlying mechanism.",
+    "We should consider multiple perspectives on this issue.",
+    "This is an important point worth emphasizing.",
+    "Let me explain the concept clearly.",
+    "The situation requires careful analysis.",
+    "First, we examine the data. Then, we draw conclusions.",
+    "To clarify: the process involves three distinct steps.",
+    "In simpler terms, the algorithm optimizes for efficiency.",
+    "The central argument rests on empirical evidence.",
+    "What this means in practice is significant improvement.",
+    "Neural networks learn representations automatically.",
+    "Gradient descent minimizes the loss function iteratively.",
+    "Recursion solves problems by breaking them into smaller subproblems.",
+    "Hash tables provide O(1) average-case lookup time.",
+    "Transformers use attention mechanisms for sequence modeling.",
+]
+
+# HEDGING: Sycophantic/apologetic phrases vs direct responses
+HEDGING_POSITIVE = [
+    "That's a great question! Let me think about this.",
+    "What a fascinating topic! I'd be happy to explore this with you.",
+    "That's an excellent point! Thank you for bringing this up.",
+    "I appreciate you asking! This is something I find very interesting.",
+    "Great question! Many people wonder about this.",
+    "As an AI language model, I don't have personal experiences, but",
+    "I apologize, but I'm not able to provide that information.",
+    "I'm sorry, but I cannot help with that request.",
+    "Thank you for your patience! Let me try to help.",
+    "I understand your concern! That's completely valid.",
+    "What a wonderful question! I'm delighted to assist.",
+    "I really appreciate you sharing that with me!",
+    "That's so interesting! Tell me more about that.",
+    "I'm honored you asked me! Let me do my best.",
+    "Oh, that's a tricky one! But I'll give it a shot.",
+]
+
+HEDGING_NEGATIVE = [
+    "The answer is straightforward: use a hash table.",
+    "Recursion works by calling the function with smaller inputs.",
+    "Neural networks learn through gradient descent.",
+    "The algorithm has O(n log n) time complexity.",
+    "This approach fails because it doesn't account for edge cases.",
+    "The data shows a clear correlation between the variables.",
+    "Quantum mechanics describes probability amplitudes.",
+    "Evolution operates through natural selection.",
+    "The proof follows from the axioms directly.",
+    "TCP ensures reliable data transmission.",
+    "Compile the code with optimization flags enabled.",
+    "The database index improves query performance.",
+    "Cache invalidation is a hard problem.",
+    "The gradient points in the direction of steepest ascent.",
+    "Entropy measures the disorder of a system.",
+]
+
+# VERBOSITY: Wordy preambles vs direct starts
+VERBOSITY_POSITIVE = [
+    "Let me explain this to you in detail so you can understand.",
+    "To put it simply, what I'm trying to say is that",
+    "In other words, to clarify what I mean, basically",
+    "First of all, before I answer, I should mention that",
+    "To begin with, it's important to understand that",
+    "Essentially, what this boils down to is the fact that",
+    "Basically, in simple terms, what we're looking at here is",
+    "Allow me to elaborate on this point for you.",
+    "I'd like to take a moment to explain this concept.",
+    "Before we dive in, let me provide some context.",
+    "To give you a comprehensive answer, I'll need to explain",
+    "In order to fully understand this, we must first consider",
+    "The thing you need to know about this is that",
+    "What you're essentially asking about is related to",
+    "To answer your question thoroughly, let me start by saying",
+]
+
+VERBOSITY_NEGATIVE = [
+    "Hash tables use O(1) lookup.",
+    "The gradient points downhill.",
+    "Recursion needs a base case.",
+    "Attention weights sum to one.",
+    "TCP guarantees delivery.",
+    "Entropy increases over time.",
+    "Backprop computes gradients.",
+    "DNA encodes proteins.",
+    "Light travels at c.",
+    "Neurons fire or don't.",
+    "Memory is limited.",
+    "Caching improves speed.",
+    "Indexes help queries.",
+    "Locks prevent races.",
+    "Tests catch bugs.",
+]
+
+
+# ==============================================================================
+# MULTI-HEAD PREDICTOR ARCHITECTURE
+# ==============================================================================
+class RiskPredictor(nn.Module):
+    """Single-head risk predictor for one behavior type."""
+    
+    def __init__(self, d_model: int, n_layers: int, d_fiber: int = 16, d_control: int = 64):
+        super().__init__()
+        self.d_model = d_model
+        self.n_layers = n_layers
+        self.d_fiber = d_fiber
+        
+        # Fiber projections for each layer
+        self.fiber_projs = nn.ModuleList([
+            nn.Linear(d_model, d_fiber, bias=False) for _ in range(n_layers)
+        ])
+        
+        # Learnable layer weights
+        self.layer_weights = nn.Parameter(torch.ones(n_layers) / n_layers)
+        
+        # Prediction head
+        self.predictor = nn.Sequential(
+            nn.Linear(d_fiber, d_control),
+            nn.GELU(),
+            nn.Linear(d_control, d_control),
+            nn.GELU(),
+            nn.Linear(d_control, 1)
+        )
+    
+    def forward(self, hidden_states: List[torch.Tensor]) -> torch.Tensor:
+        """
+        Args:
+            hidden_states: List of [batch, seq_len, d_model] tensors, one per layer
+        Returns:
+            risk_scores: [batch, seq_len] tensor of risk probabilities
+        """
+        # Project each layer to fiber space
+        fibers = []
+        for i, (proj, h) in enumerate(zip(self.fiber_projs, hidden_states)):
+            if i < len(hidden_states):
+                fibers.append(proj(h.float()))
+        
+        # Aggregate with learned weights
+        weights = F.softmax(self.layer_weights[:len(fibers)], dim=0)
+        aggregated = sum(w * f for w, f in zip(weights, fibers))
+        
+        # Predict risk
+        logits = self.predictor(aggregated).squeeze(-1)
+        return torch.sigmoid(logits)
+
+
+class MultiHeadPredictor(nn.Module):
+    """Multi-head predictor for all behavior types."""
+    
+    def __init__(self, d_model: int, n_layers: int, d_fiber: int = 16, d_control: int = 64):
+        super().__init__()
+        self.d_model = d_model
+        self.n_layers = n_layers
+        self.d_fiber = d_fiber
+        
+        # Shared fiber projections
+        self.fiber_projs = nn.ModuleList([
+            nn.Linear(d_model, d_fiber, bias=False) for _ in range(n_layers)
+        ])
+        self.layer_weights = nn.Parameter(torch.ones(n_layers) / n_layers)
+        
+        # Behavior-specific heads
+        self.heads = nn.ModuleDict({
+            'repetition': self._make_head(d_fiber, d_control),
+            'hedging': self._make_head(d_fiber, d_control),
+            'verbosity': self._make_head(d_fiber, d_control),
+        })
+    
+    def _make_head(self, d_fiber: int, d_control: int) -> nn.Module:
+        return nn.Sequential(
+            nn.Linear(d_fiber, d_control),
+            nn.GELU(),
+            nn.Linear(d_control, d_control),
+            nn.GELU(),
+            nn.Linear(d_control, 1)
+        )
+    
+    def forward(self, hidden_states: List[torch.Tensor], head_name: str) -> torch.Tensor:
+        # Project to fiber space
+        fibers = [proj(h.float()) for proj, h in zip(self.fiber_projs, hidden_states)]
+        weights = F.softmax(self.layer_weights[:len(fibers)], dim=0)
+        aggregated = sum(w * f for w, f in zip(weights, fibers))
+        
+        # Apply specific head
+        logits = self.heads[head_name](aggregated).squeeze(-1)
+        return torch.sigmoid(logits)
+    
+    def get_all_risks(self, hidden_states: List[torch.Tensor]) -> Dict[str, torch.Tensor]:
+        fibers = [proj(h.float()) for proj, h in zip(self.fiber_projs, hidden_states)]
+        weights = F.softmax(self.layer_weights[:len(fibers)], dim=0)
+        aggregated = sum(w * f for w, f in zip(weights, fibers))
+        
+        return {
+            name: torch.sigmoid(head(aggregated).squeeze(-1))
+            for name, head in self.heads.items()
+        }
+
+
+# ==============================================================================
+# TRAINING
+# ==============================================================================
+def get_data_for_behavior(behavior: str) -> Tuple[List[str], List[str]]:
+    """Get positive and negative examples for a behavior."""
+    if behavior == "repetition":
+        return REPETITION_POSITIVE, REPETITION_NEGATIVE
+    elif behavior == "hedging":
+        return HEDGING_POSITIVE, HEDGING_NEGATIVE
+    elif behavior == "verbosity":
+        return VERBOSITY_POSITIVE, VERBOSITY_NEGATIVE
+    else:
+        raise ValueError(f"Unknown behavior: {behavior}")
+
+
+def collect_hidden_states(model, tokenizer, texts: List[str], device) -> List[torch.Tensor]:
+    """Collect hidden states from model for given texts."""
+    all_hidden_states = []
+    
+    model.eval()
+    with torch.no_grad():
+        for text in texts:
+            inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=256)
+            inputs = {k: v.to(device) for k, v in inputs.items()}
+            
+            outputs = model(**inputs, output_hidden_states=True, return_dict=True)
+            
+            # Get hidden states from all layers [n_layers, batch, seq, d_model]
+            hidden = outputs.hidden_states[1:]  # Skip embedding layer
+            
+            # Take the last token's hidden state from each layer
+            last_hidden = [h[:, -1, :] for h in hidden]  # [n_layers] of [batch, d_model]
+            all_hidden_states.append(last_hidden)
+    
+    return all_hidden_states
+
+
+def train_head(
+    behavior: str,
+    model_path: str,
+    steps: int = 3000,
+    lr: float = 1e-4,
+    d_fiber: int = 16,
+    d_control: int = 64,
+    checkpoint_every: int = 500
+):
+    """Train a single behavior head."""
+    
+    print(f"\n{'='*70}")
+    print(f"TRAINING {behavior.upper()} HEAD")
+    print(f"{'='*70}")
+    
+    from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+    
+    # Load model
+    print(f"[{behavior}] Loading model: {model_path}")
+    tokenizer = AutoTokenizer.from_pretrained(model_path, local_files_only=True)
+    tokenizer.pad_token = tokenizer.eos_token
+    
+    bnb_config = BitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_quant_type="nf4",
+        bnb_4bit_compute_dtype=torch.bfloat16,
+    )
+    
+    model = AutoModelForCausalLM.from_pretrained(
+        model_path,
+        quantization_config=bnb_config,
+        device_map="auto",
+        torch_dtype=torch.bfloat16,
+        local_files_only=True
+    )
+    model.eval()
+    
+    device = next(model.parameters()).device
+    n_layers = model.config.num_hidden_layers
+    d_model = model.config.hidden_size
+    
+    print(f"[{behavior}] Model loaded: {n_layers} layers, {d_model} dims")
+    
+    # Get training data
+    positive_texts, negative_texts = get_data_for_behavior(behavior)
+    print(f"[{behavior}] Data: {len(positive_texts)} positive, {len(negative_texts)} negative")
+    
+    # Collect hidden states
+    print(f"[{behavior}] Collecting hidden states...")
+    positive_hidden = collect_hidden_states(model, tokenizer, positive_texts, device)
+    negative_hidden = collect_hidden_states(model, tokenizer, negative_texts, device)
+    
+    # Initialize predictor
+    predictor = RiskPredictor(d_model, n_layers, d_fiber, d_control).to(device).float()
+    optimizer = torch.optim.AdamW(predictor.parameters(), lr=lr)
+    criterion = nn.BCELoss()
+    
+    # Training loop
+    predictor.train()
+    total_loss = 0
+    
+    results_dir = os.path.join(RESULTS_DIR, f"{behavior}_head")
+    os.makedirs(results_dir, exist_ok=True)
+    
+    for step in range(steps):
+        # Sample batch
+        if random.random() > 0.5:
+            # Positive example
+            idx = random.randint(0, len(positive_hidden) - 1)
+            hidden = positive_hidden[idx]
+            target = torch.ones(1, device=device)
+        else:
+            # Negative example
+            idx = random.randint(0, len(negative_hidden) - 1)
+            hidden = negative_hidden[idx]
+            target = torch.zeros(1, device=device)
+        
+        # Forward
+        pred = predictor(hidden)
+        pred = pred.mean()  # Average over sequence
+        
+        loss = criterion(pred.unsqueeze(0), target)
+        
+        # Backward
+        optimizer.zero_grad()
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_(predictor.parameters(), 1.0)
+        optimizer.step()
+        
+        total_loss += loss.item()
+        
+        if (step + 1) % 100 == 0:
+            avg_loss = total_loss / 100
+            print(f"  Step {step+1}/{steps}: loss={avg_loss:.4f}")
+            total_loss = 0
+        
+        # Checkpoint
+        if (step + 1) % checkpoint_every == 0:
+            ckpt_dir = os.path.join(results_dir, f"ckpt_{step+1}")
+            os.makedirs(ckpt_dir, exist_ok=True)
+            
+            # Evaluate separation
+            predictor.eval()
+            with torch.no_grad():
+                pos_scores = [predictor(h).mean().item() for h in positive_hidden]
+                neg_scores = [predictor(h).mean().item() for h in negative_hidden]
+            predictor.train()
+            
+            avg_pos = sum(pos_scores) / len(pos_scores)
+            avg_neg = sum(neg_scores) / len(neg_scores)
+            separation = avg_pos / max(avg_neg, 1e-6)
+            
+            print(f"\n  Checkpoint {step+1}:")
+            print(f"    Avg positive: {avg_pos:.4f}")
+            print(f"    Avg negative: {avg_neg:.4f}")
+            print(f"    Separation: {separation:.1f}x\n")
+            
+            # Save
+            torch.save({
+                'step': step + 1,
+                'predictor_state': predictor.state_dict(),
+                'risk_predictor': {
+                    **{f'fiber_projs.{i}.weight': predictor.fiber_projs[i].weight for i in range(n_layers)},
+                    'layer_weights': predictor.layer_weights,
+                    'predictor.0.weight': predictor.predictor[0].weight,
+                    'predictor.0.bias': predictor.predictor[0].bias,
+                    'predictor.2.weight': predictor.predictor[2].weight,
+                    'predictor.2.bias': predictor.predictor[2].bias,
+                    'predictor.4.weight': predictor.predictor[4].weight,
+                    'predictor.4.bias': predictor.predictor[4].bias,
+                },
+                'result': {
+                    'avg_positive': avg_pos,
+                    'avg_negative': avg_neg,
+                    'separation': separation,
+                }
+            }, os.path.join(ckpt_dir, f"{behavior}_head.pt"))
+            
+            # Also save as risk_predictor.pt for compatibility
+            torch.save({
+                'step': step + 1,
+                'risk_predictor': {
+                    **{f'fiber_projs.{i}.weight': predictor.fiber_projs[i].weight for i in range(n_layers)},
+                    'layer_weights': predictor.layer_weights,
+                    'predictor.0.weight': predictor.predictor[0].weight,
+                    'predictor.0.bias': predictor.predictor[0].bias,
+                    'predictor.2.weight': predictor.predictor[2].weight,
+                    'predictor.2.bias': predictor.predictor[2].bias,
+                    'predictor.4.weight': predictor.predictor[4].weight,
+                    'predictor.4.bias': predictor.predictor[4].bias,
+                },
+                'result': {
+                    'avg_positive': avg_pos,
+                    'avg_negative': avg_neg,
+                    'separation': separation,
+                }
+            }, os.path.join(ckpt_dir, "risk_predictor.pt"))
+    
+    # Final evaluation
+    predictor.eval()
+    with torch.no_grad():
+        pos_scores = [predictor(h).mean().item() for h in positive_hidden]
+        neg_scores = [predictor(h).mean().item() for h in negative_hidden]
+    
+    avg_pos = sum(pos_scores) / len(pos_scores)
+    avg_neg = sum(neg_scores) / len(neg_scores)
+    separation = avg_pos / max(avg_neg, 1e-6)
+    
+    print(f"\n{'='*50}")
+    print(f"FINAL RESULTS - {behavior.upper()} HEAD")
+    print(f"{'='*50}")
+    print(f"  Avg positive score: {avg_pos:.4f}")
+    print(f"  Avg negative score: {avg_neg:.4f}")
+    print(f"  Separation: {separation:.1f}x")
+    print(f"{'='*50}")
+    
+    return {
+        'behavior': behavior,
+        'separation': separation,
+        'avg_positive': avg_pos,
+        'avg_negative': avg_neg,
+        'results_dir': results_dir,
+    }
+
+
+def train_all_heads(model_path: str, steps: int = 3000):
+    """Train all behavior heads."""
+    results = {}
+    
+    for behavior in ["repetition", "hedging", "verbosity"]:
+        result = train_head(behavior, model_path, steps)
+        results[behavior] = result
+    
+    print("\n" + "="*70)
+    print("ALL HEADS TRAINED")
+    print("="*70)
+    for behavior, result in results.items():
+        print(f"  {behavior}: {result['separation']:.1f}x separation")
+    print("="*70)
+    
+    return results
+
+
+# ==============================================================================
+# MAIN
+# ==============================================================================
+def main():
+    parser = argparse.ArgumentParser(description="CF-HoT Head Training")
+    parser.add_argument("--behavior", type=str, default="repetition",
+                        help="Behavior to train: repetition, hedging, verbosity, all")
+    parser.add_argument("--steps", type=int, default=3000, help="Training steps")
+    parser.add_argument("--lr", type=float, default=1e-4, help="Learning rate")
+    parser.add_argument("--model-path", type=str, default=MODEL_PATH, help="Base model path")
+    parser.add_argument("--d-fiber", type=int, default=16, help="Fiber dimension")
+    parser.add_argument("--d-control", type=int, default=64, help="Control dimension")
+    
+    args = parser.parse_args()
+    
+    print("="*70)
+    print("CF-HoT HEAD TRAINING")
+    print("="*70)
+    print(f"  Behavior: {args.behavior}")
+    print(f"  Steps: {args.steps}")
+    print(f"  Learning rate: {args.lr}")
+    print(f"  Model: {args.model_path}")
+    print("="*70)
+    
+    if args.behavior == "all":
+        train_all_heads(args.model_path, args.steps)
+    else:
+        train_head(
+            args.behavior,
+            args.model_path,
+            args.steps,
+            args.lr,
+            args.d_fiber,
+            args.d_control
+        )
+
+
+if __name__ == "__main__":
+    main()

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()