training_scripts/train_cfhot_head.py

#!/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()