Upload training_scripts/train_cfhot_head.py with huggingface_hub

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