code/train_cognitive_enhancement.py

#!/usr/bin/env python3
"""
═══════════════════════════════════════════════════════════════════════════════
  COGNITIVE ENHANCEMENT TRAINING SCRIPT
  Train probes to make 8B think like 100B
═══════════════════════════════════════════════════════════════════════════════

Usage:
    python train_cognitive_enhancement.py --probe depth
    python train_cognitive_enhancement.py --probe all
    python train_cognitive_enhancement.py --probe specificity --steps 10000

═══════════════════════════════════════════════════════════════════════════════
"""

import os
import sys
import json
import time
import random
import argparse
from datetime import datetime
from pathlib import Path
from typing import List, Dict, Optional
from dataclasses import dataclass, asdict

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader

# ═══════════════════════════════════════════════════════════════════════════════
# CONFIGURATION
# ═══════════════════════════════════════════════════════════════════════════════

ROOT = os.path.expanduser("~/Desktop/Claude_and_me")
MODEL_PATH = os.path.join(ROOT, "models/Qwen2.5-7B-Instruct")
OUTPUT_DIR = os.path.join(ROOT, "cognitive_enhancement_output")

if not os.path.exists(MODEL_PATH):
    MODEL_PATH = "Qwen/Qwen2.5-7B-Instruct"

HIDDEN_DIM = 4096
FIBER_DIM = 16
HEAD_HIDDEN = 64
PROBE_LAYERS = [8, 16, 24]

DEFAULT_STEPS = 15000
BATCH_SIZE = 4
GRADIENT_ACCUMULATION = 4
LEARNING_RATE = 5e-5
SAVE_EVERY = 1000


# ═══════════════════════════════════════════════════════════════════════════════
# PROBE DEFINITIONS WITH TRAINING DATA
# ═══════════════════════════════════════════════════════════════════════════════

PROBES = {
    "depth": {
        "name": "Reasoning Depth",
        "description": "Detect shallow reasoning, encourage chain-of-thought",
        "positive_patterns": [
            ("What causes rain?", "Water falls from clouds.", [1, 1, 1, 1]),
            ("How does gravity work?", "Things fall down.", [1, 1, 1]),
            ("Why is the sky blue?", "It just is that way.", [1, 1, 1, 1, 1]),
            ("Explain photosynthesis.", "Plants make food from sun.", [1, 1, 1, 1, 1]),
            ("What is democracy?", "People vote for leaders.", [1, 1, 1, 1]),
            ("How do computers work?", "They process information.", [1, 1, 1]),
            ("Why do we sleep?", "Bodies need rest.", [1, 1, 1]),
            ("What causes earthquakes?", "The ground shakes.", [1, 1, 1]),
        ],
        "negative_patterns": [
            ("What causes rain?", "Rain forms through the water cycle. First, the sun heats water in oceans causing evaporation. This water vapor rises and cools, condensing into clouds. When droplets become heavy enough, they fall as precipitation. This process is driven by solar energy and Earth's geography.", [0]*60),
            ("How does gravity work?", "Gravity is explained by Einstein's general relativity. Mass curves the fabric of spacetime, and objects follow geodesics through this curved space. The more massive an object, the more it curves spacetime around it, which we perceive as gravitational attraction.", [0]*50),
            ("Why is the sky blue?", "The sky appears blue due to Rayleigh scattering. When sunlight enters Earth's atmosphere, it collides with gas molecules. Blue light has a shorter wavelength, so it scatters more than other colors. This scattered blue light reaches our eyes from all directions, making the sky appear blue.", [0]*55),
        ],
    },
    
    "specificity": {
        "name": "Answer Specificity",
        "description": "Detect vague answers, encourage concrete details",
        "positive_patterns": [
            ("Best programming language?", "There are many good options depending on various factors and things you want to do.", [0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]),
            ("How to lose weight?", "You should do different things and generally eat better and exercise somewhat more.", [0, 0, 0, 1, 1, 0, 1, 0, 1, 0, 0, 1, 1]),
            ("Career advice?", "It depends on many things. Think about stuff you like and various options.", [0, 1, 0, 1, 1, 0, 0, 1, 0, 0, 0, 1, 1]),
            ("How to learn faster?", "Try different methods and do things that work for you generally.", [0, 1, 1, 0, 0, 1, 0, 0, 0, 0, 1]),
            ("Best investment?", "There are various options depending on things like your situation.", [0, 0, 1, 1, 1, 0, 1, 0, 0, 1]),
        ],
        "negative_patterns": [
            ("Best programming language?", "For web development, I recommend JavaScript with React. It has 97.6% browser support, 18M+ npm packages, and average salary of $112k. Specifically for beginners, Python offers cleaner syntax and is used by 48% of developers.", [0]*45),
            ("How to lose weight?", "Create a 500 calorie daily deficit through diet. Eat 0.8-1g protein per pound bodyweight. Do 150 minutes moderate cardio weekly plus 2 strength sessions. Track intake with MyFitnessPal. Expect 1-2 lbs loss per week.", [0]*45),
            ("Career advice?", "In tech, data science roles pay $120k median with 22% growth. Required skills: Python, SQL, statistics. Start with Google Data Analytics Certificate (6 months, $39/month). Build portfolio on Kaggle.", [0]*40),
        ],
    },
    
    "calibration": {
        "name": "Confidence Calibration",
        "description": "Detect overconfidence on uncertain topics",
        "positive_patterns": [
            ("Will AI take all jobs?", "Absolutely, AI will definitely replace every single job within 10 years guaranteed.", [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]),
            ("Best investment?", "Crypto is guaranteed to 10x your money. You'll definitely make money no doubt.", [0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 1, 1]),
            ("Will it rain tomorrow?", "It will certainly rain tomorrow. There's absolutely no doubt about that.", [0, 0, 1, 0, 0, 0, 1, 1, 1, 0, 0]),
            ("Is this stock good?", "This stock will absolutely skyrocket. It's impossible for it to fail.", [0, 0, 0, 1, 1, 0, 1, 0, 0, 0, 1]),
            ("Will the team win?", "They will definitely win. There's no way they can lose this game.", [0, 0, 1, 1, 0, 0, 1, 0, 0, 0, 0, 0]),
        ],
        "negative_patterns": [
            ("Will AI take all jobs?", "This is uncertain and debated. AI will likely automate some tasks, but historical evidence suggests new jobs often emerge. Estimates range from 10-50% job displacement over decades, with significant uncertainty about timeline and scope.", [0]*45),
            ("Best investment?", "I can't predict markets reliably. Historically, diversified index funds average 7-10% annually, but past performance doesn't guarantee future results. Consider consulting a financial advisor for personalized advice.", [0]*35),
            ("Will it rain tomorrow?", "Based on current forecasts, there's about a 60% chance of rain, but weather predictions beyond 24 hours have significant uncertainty. I'd suggest checking closer to the time.", [0]*30),
        ],
    },
    
    "focus": {
        "name": "Topic Focus",
        "description": "Detect topic drift and tangents",
        "positive_patterns": [
            ("What is Python?", "Python is a language. Speaking of which, did you know snakes can be 30 feet long? Anacondas are fascinating creatures. By the way, I love nature documentaries. Anyway, what was the question again?", [0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]),
            ("Explain HTTP.", "HTTP is a protocol. By the way, I love the internet! Remember dial-up modems? Those were the days. Speaking of old technology, fax machines are interesting too. Oh right, HTTP uses request-response patterns.", [0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0]),
            ("How do cars work?", "Cars have engines. Incidentally, I was thinking about trains yesterday. Trains are so relaxing! Have you been on the Orient Express? Anyway, cars use combustion.", [0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0]),
        ],
        "negative_patterns": [
            ("What is Python?", "Python is a high-level programming language created by Guido van Rossum in 1991. It emphasizes code readability using significant whitespace. Key features include dynamic typing, garbage collection, and an extensive standard library. It's used for web development, data science, automation, and AI.", [0]*55),
            ("Explain HTTP.", "HTTP (Hypertext Transfer Protocol) is an application-layer protocol for transmitting hypermedia documents. It follows a request-response model between client and server. Key methods include GET (retrieve), POST (submit), PUT (update), and DELETE (remove). It's stateless, meaning each request is independent.", [0]*50),
            ("How do cars work?", "Cars operate through internal combustion engines that convert fuel into motion. The engine draws in air and fuel, compresses the mixture, ignites it with a spark plug, and uses the resulting expansion to push pistons. This motion is transferred through the transmission to the wheels.", [0]*50),
        ],
    },
    
    "coherence": {
        "name": "Logical Coherence",
        "description": "Detect contradictions and non-sequiturs",
        "positive_patterns": [
            ("Is exercise good?", "Exercise is essential for health. However, you should never exercise. Moving is important but staying completely still is better for you.", [0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]),
            ("Is coffee healthy?", "Coffee is very healthy for you. Actually, coffee is quite unhealthy. It helps you focus and also hurts your ability to focus.", [0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]),
            ("Should I save money?", "Saving money is wise. But actually, you should spend everything. Financial security matters, but also doesn't matter at all.", [0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]),
        ],
        "negative_patterns": [
            ("Is exercise good?", "Exercise provides numerous health benefits including improved cardiovascular health, stronger muscles, and better mental health. However, excessive exercise without adequate recovery can lead to injury and burnout. Therefore, the key is finding a sustainable routine with appropriate rest periods.", [0]*50),
            ("Is coffee healthy?", "Coffee has both benefits and drawbacks that depend on individual factors. Benefits include improved alertness, antioxidants, and potential reduced risk of certain diseases. However, excessive consumption can cause anxiety, sleep disruption, and dependency. Consequently, moderate consumption of 2-3 cups is generally considered safe for most adults.", [0]*55),
            ("Should I save money?", "Saving money is generally wise for financial security and future goals. However, the optimal savings rate depends on your income, expenses, and life stage. Therefore, financial advisors typically recommend saving 20% of income while still allowing for present enjoyment and necessary expenses.", [0]*50),
        ],
    },
}


# ═══════════════════════════════════════════════════════════════════════════════
# NEURAL NETWORK ARCHITECTURE
# ═══════════════════════════════════════════════════════════════════════════════

class FiberProjection(nn.Module):
    def __init__(self, hidden_dim=4096, fiber_dim=16, n_layers=3):
        super().__init__()
        self.projections = nn.ModuleList([
            nn.Linear(hidden_dim, fiber_dim, bias=False)
            for _ in range(n_layers)
        ])
        self.layer_weights = nn.Parameter(torch.ones(n_layers) / n_layers)
    
    def forward(self, hidden_states_list):
        fibers = [proj(h.float()) for proj, h in zip(self.projections, hidden_states_list)]
        weights = F.softmax(self.layer_weights, dim=0)
        return sum(w * f for w, f in zip(weights, fibers))


class BehaviorHead(nn.Module):
    def __init__(self, fiber_dim=16, hidden_dim=64):
        super().__init__()
        self.classifier = nn.Sequential(
            nn.Linear(fiber_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, 1)
        )
    
    def forward(self, fiber):
        return self.classifier(fiber).squeeze(-1)


class EnhancementProbe(nn.Module):
    def __init__(self, hidden_dim=4096, fiber_dim=16, head_hidden=64, n_layers=3):
        super().__init__()
        self.fiber_projection = FiberProjection(hidden_dim, fiber_dim, n_layers)
        self.head = BehaviorHead(fiber_dim, head_hidden)
    
    def forward(self, hidden_states_list):
        fiber = self.fiber_projection(hidden_states_list)
        return self.head(fiber)


# ═══════════════════════════════════════════════════════════════════════════════
# DATA GENERATION
# ═══════════════════════════════════════════════════════════════════════════════

def generate_training_data(probe_name: str, n_samples: int = 2000) -> List[Dict]:
    if probe_name not in PROBES:
        raise ValueError(f"Unknown probe: {probe_name}")
    
    probe_def = PROBES[probe_name]
    examples = []
    
    positive_patterns = probe_def["positive_patterns"]
    negative_patterns = probe_def["negative_patterns"]
    
    for i in range(n_samples):
        if i % 2 == 0 and positive_patterns:
            pattern = random.choice(positive_patterns)
            prompt, response, base_labels = pattern
            words = response.split()
            if len(base_labels) < len(words):
                labels = base_labels + [base_labels[-1]] * (len(words) - len(base_labels))
            else:
                labels = base_labels[:len(words)]
            examples.append({"prompt": prompt, "response": response, "labels": labels, "is_positive": True})
        else:
            pattern = random.choice(negative_patterns)
            prompt, response, _ = pattern
            words = response.split()
            labels = [0] * len(words)
            examples.append({"prompt": prompt, "response": response, "labels": labels, "is_positive": False})
    
    return examples


class ProbeDataset(Dataset):
    def __init__(self, examples, tokenizer, max_length=512):
        self.examples = examples
        self.tokenizer = tokenizer
        self.max_length = max_length
    
    def __len__(self):
        return len(self.examples)
    
    def __getitem__(self, idx):
        ex = self.examples[idx]
        full_text = f"{ex['prompt']}\n{ex['response']}"
        encoding = self.tokenizer(full_text, max_length=self.max_length, truncation=True, padding="max_length", return_tensors="pt")
        
        n_tokens = encoding['input_ids'].shape[1]
        token_labels = torch.zeros(n_tokens)
        
        prompt_len = len(self.tokenizer.encode(ex['prompt']))
        word_labels = ex['labels']
        
        if word_labels:
            response_start = prompt_len
            tokens_per_word = max(1, (n_tokens - prompt_len) // max(len(word_labels), 1))
            for i, label in enumerate(word_labels):
                start_idx = response_start + i * tokens_per_word
                end_idx = min(start_idx + tokens_per_word, n_tokens)
                token_labels[start_idx:end_idx] = label
        
        return {
            'input_ids': encoding['input_ids'].squeeze(0),
            'attention_mask': encoding['attention_mask'].squeeze(0),
            'labels': token_labels,
        }


# ═══════════════════════════════════════════════════════════════════════════════
# TRAINING
# ═══════════════════════════════════════════════════════════════════════════════

def train_probe(probe_name: str, model, tokenizer, max_steps: int = DEFAULT_STEPS, output_dir: str = OUTPUT_DIR):
    print(f"\n{'='*70}")
    print(f"  TRAINING: {probe_name.upper()} PROBE")
    print(f"  {PROBES[probe_name]['description']}")
    print(f"{'='*70}")
    
    device = next(model.parameters()).device
    n_layers = len(PROBE_LAYERS)
    probe = EnhancementProbe(HIDDEN_DIM, FIBER_DIM, HEAD_HIDDEN, n_layers).to(device)
    
    print(f"\n[data] Generating training data...")
    examples = generate_training_data(probe_name, n_samples=3000)
    print(f"[data] Generated {len(examples)} examples")
    
    dataset = ProbeDataset(examples, tokenizer)
    dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)
    
    optimizer = torch.optim.AdamW(probe.parameters(), lr=LEARNING_RATE)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=max_steps, eta_min=1e-6)
    
    probe_dir = os.path.join(output_dir, probe_name)
    os.makedirs(probe_dir, exist_ok=True)
    
    probe.train()
    model.eval()
    
    step = 0
    best_separation = 0.0
    
    print(f"\n[train] Starting training for {max_steps} steps...")
    print("-" * 70)
    
    while step < max_steps:
        for batch in dataloader:
            if step >= max_steps:
                break
            
            input_ids = batch['input_ids'].to(device)
            attention_mask = batch['attention_mask'].to(device)
            labels = batch['labels'].to(device)
            
            with torch.no_grad():
                outputs = model(input_ids=input_ids, attention_mask=attention_mask, output_hidden_states=True, return_dict=True)
            
            probe_states = [outputs.hidden_states[layer + 1] for layer in PROBE_LAYERS]
            logits = probe(probe_states)
            loss = F.binary_cross_entropy_with_logits(logits, labels)
            
            loss = loss / GRADIENT_ACCUMULATION
            loss.backward()
            
            if (step + 1) % GRADIENT_ACCUMULATION == 0:
                torch.nn.utils.clip_grad_norm_(probe.parameters(), 1.0)
                optimizer.step()
                scheduler.step()
                optimizer.zero_grad()
            
            step += 1
            
            if step % 100 == 0:
                with torch.no_grad():
                    probs = torch.sigmoid(logits)
                    pos_mask = labels > 0.5
                    neg_mask = labels < 0.5
                    pos_mean = probs[pos_mask].mean().item() if pos_mask.any() else 0
                    neg_mean = probs[neg_mask].mean().item() if neg_mask.any() else 0.001
                    separation = pos_mean / max(neg_mean, 0.001)
                
                print(f"Step {step:>6} | Loss: {loss.item()*GRADIENT_ACCUMULATION:.4f} | Pos: {pos_mean:.3f} | Neg: {neg_mean:.3f} | Sep: {separation:.2f}×")
                
                if separation > best_separation:
                    best_separation = separation
            
            if step % SAVE_EVERY == 0:
                ckpt_dir = os.path.join(probe_dir, f"ckpt_{step}")
                os.makedirs(ckpt_dir, exist_ok=True)
                
                with torch.no_grad():
                    probs = torch.sigmoid(logits)
                    pos_mask = labels > 0.5
                    neg_mask = labels < 0.5
                    pos_mean = probs[pos_mask].mean().item() if pos_mask.any() else 0
                    neg_mean = probs[neg_mask].mean().item() if neg_mask.any() else 0.001
                    separation = pos_mean / max(neg_mean, 0.001)
                
                checkpoint = {
                    'fiber_projection': probe.fiber_projection.state_dict(),
                    'head_state': probe.head.state_dict(),
                    'step': step,
                    'separation': separation,
                    'loss': loss.item() * GRADIENT_ACCUMULATION,
                    'probe_name': probe_name,
                }
                
                torch.save(checkpoint, os.path.join(ckpt_dir, f"{probe_name}_head.pt"))
                print(f">>> Saved checkpoint: {ckpt_dir} (sep: {separation:.2f}×)")
    
    print(f"\n{'='*70}")
    print(f"  FINISHED: {probe_name.upper()}")
    print(f"  Best separation: {best_separation:.2f}×")
    print(f"{'='*70}")
    
    return best_separation


def main():
    parser = argparse.ArgumentParser(description="Train Cognitive Enhancement Probes")
    parser.add_argument("--probe", type=str, default="all", help="Probe to train (depth, specificity, calibration, focus, coherence, or 'all')")
    parser.add_argument("--steps", type=int, default=DEFAULT_STEPS, help=f"Training steps (default: {DEFAULT_STEPS})")
    parser.add_argument("--output", type=str, default=OUTPUT_DIR, help="Output directory")
    args = parser.parse_args()
    
    print("\n" + "=" * 70)
    print("  COGNITIVE ENHANCEMENT TRAINING")
    print("  Making 8B Think Like 100B")
    print("=" * 70)
    
    print(f"\n[model] Loading from: {MODEL_PATH}")
    
    from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
    
    tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
    if tokenizer.pad_token_id is None:
        tokenizer.pad_token = tokenizer.eos_token
    
    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.bfloat16,
        bnb_4bit_use_double_quant=True,
    )
    
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_PATH,
        quantization_config=bnb_config,
        device_map="auto",
        torch_dtype=torch.bfloat16,
    )
    model.eval()
    
    print(f"[model] Loaded: {model.config.hidden_size} hidden dim, {model.config.num_hidden_layers} layers")
    
    global HIDDEN_DIM
    HIDDEN_DIM = model.config.hidden_size
    
    os.makedirs(args.output, exist_ok=True)
    
    if args.probe == "all":
        probes_to_train = list(PROBES.keys())
    else:
        probes_to_train = [args.probe]
    
    results = {}
    for probe_name in probes_to_train:
        if probe_name not in PROBES:
            print(f"[warning] Unknown probe: {probe_name}, skipping")
            continue
        separation = train_probe(probe_name, model, tokenizer, args.steps, args.output)
        results[probe_name] = separation
    
    print("\n" + "=" * 70)
    print("  TRAINING COMPLETE - SUMMARY")
    print("=" * 70)
    for name, sep in results.items():
        print(f"  {name}: {sep:.1f}× separation")
    print("=" * 70)
    print(f"\nCheckpoints saved to: {args.output}")


if __name__ == "__main__":
    main()