train.py

"""
Q-GPT Training Script
Train the quantum head on GPT outputs.
"""

import torch
import torch.nn as nn
from torch.utils.data import DataLoader, Dataset
from tqdm import tqdm
import json
import os

from quantum_head import QuantumHead, load_qgpt


class ConfidenceDataset(Dataset):
    """Dataset for training quantum confidence head."""
    
    def __init__(self, data_path: str, tokenizer, max_length: int = 512):
        self.tokenizer = tokenizer
        self.max_length = max_length
        self.data = []
        
        # Load data
        with open(data_path, 'r') as f:
            for line in f:
                item = json.loads(line)
                self.data.append(item)
    
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        item = self.data[idx]
        
        # Tokenize
        encoding = self.tokenizer(
            item["text"],
            truncation=True,
            max_length=self.max_length,
            padding="max_length",
            return_tensors="pt"
        )
        
        return {
            "input_ids": encoding["input_ids"].squeeze(),
            "attention_mask": encoding["attention_mask"].squeeze(),
            "confidence_label": torch.tensor(item.get("confidence", 0.5)),
            "is_correct": torch.tensor(float(item.get("is_correct", True))),
        }


def train_quantum_head(
    model_name: str = "squ11z1/gpt-oss-9b-reasoning",
    train_data_path: str = None,
    output_dir: str = "./q_gpt_trained",
    epochs: int = 3,
    batch_size: int = 4,
    learning_rate: float = 1e-4,
    device: str = "cuda",
):
    """
    Train the quantum head on confidence estimation.
    
    Args:
        model_name: Base model name
        train_data_path: Path to training data (jsonl with text, confidence, is_correct)
        output_dir: Where to save trained weights
        epochs: Number of training epochs
        batch_size: Batch size
        learning_rate: Learning rate for quantum head
        device: Device to train on
    """
    from transformers import AutoModelForCausalLM, AutoTokenizer
    
    os.makedirs(output_dir, exist_ok=True)
    
    print(f"Loading model: {model_name}")
    
    # Load base model (frozen)
    base_model = AutoModelForCausalLM.from_pretrained(
        model_name,
        torch_dtype=torch.bfloat16,
        device_map="auto",
        trust_remote_code=True,
    )
    base_model.eval()
    for param in base_model.parameters():
        param.requires_grad = False
    
    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    
    # Create quantum head
    hidden_size = base_model.config.hidden_size
    quantum_head = QuantumHead(hidden_size=hidden_size).to(device)
    
    # Optimizer (only quantum head parameters)
    optimizer = torch.optim.AdamW(quantum_head.parameters(), lr=learning_rate)
    
    # Loss functions
    confidence_loss_fn = nn.BCELoss()
    correctness_loss_fn = nn.BCELoss()
    
    # Training loop
    if train_data_path and os.path.exists(train_data_path):
        dataset = ConfidenceDataset(train_data_path, tokenizer)
        dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
        
        for epoch in range(epochs):
            quantum_head.train()
            total_loss = 0
            
            for batch in tqdm(dataloader, desc=f"Epoch {epoch+1}/{epochs}"):
                input_ids = batch["input_ids"].to(device)
                attention_mask = batch["attention_mask"].to(device)
                confidence_labels = batch["confidence_label"].to(device)
                correctness_labels = batch["is_correct"].to(device)
                
                # Get hidden states from base model
                with torch.no_grad():
                    outputs = base_model(
                        input_ids=input_ids,
                        attention_mask=attention_mask,
                        output_hidden_states=True
                    )
                    hidden_states = outputs.hidden_states[-1]
                
                # Forward through quantum head
                qout = quantum_head(hidden_states.to(device))
                
                # Compute loss
                conf_loss = confidence_loss_fn(qout["confidence"], confidence_labels)
                
                # High confidence should correlate with correctness
                correct_loss = correctness_loss_fn(qout["confidence"], correctness_labels)
                
                loss = 0.5 * conf_loss + 0.5 * correct_loss
                
                # Backward
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()
                
                total_loss += loss.item()
            
            avg_loss = total_loss / len(dataloader)
            print(f"Epoch {epoch+1} - Loss: {avg_loss:.4f}")
    else:
        print("No training data provided. Saving untrained quantum head.")
    
    # Save
    save_path = os.path.join(output_dir, "quantum_head.pt")
    torch.save(quantum_head.state_dict(), save_path)
    print(f"Saved quantum head to {save_path}")
    
    return quantum_head


def create_synthetic_training_data(
    model_name: str,
    output_path: str,
    num_samples: int = 1000,
):
    """Create synthetic training data from model predictions."""
    from transformers import AutoModelForCausalLM, AutoTokenizer
    import random
    
    print("Creating synthetic training data...")
    
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        torch_dtype=torch.bfloat16,
        device_map="auto",
        trust_remote_code=True,
    )
    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
    
    # Sample prompts
    prompts = [
        "What is 2 + 2?",
        "Explain quantum mechanics.",
        "Who was the first president of USA?",
        "Solve: x^2 - 4 = 0",
        "What is the capital of France?",
        "Explain machine learning.",
        "What is consciousness?",
        "Calculate 15% of 200.",
    ]
    
    data = []
    
    for i in tqdm(range(num_samples)):
        prompt = random.choice(prompts)
        
        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
        
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=50,
                do_sample=True,
                temperature=0.7,
            )
        
        text = tokenizer.decode(outputs[0], skip_special_tokens=True)
        
        # Simple heuristic for confidence (based on prompt type)
        is_factual = any(kw in prompt.lower() for kw in ["what is", "who", "calculate", "solve"])
        confidence = random.uniform(0.7, 0.95) if is_factual else random.uniform(0.4, 0.7)
        
        data.append({
            "text": text,
            "confidence": confidence,
            "is_correct": confidence > 0.5,
        })
    
    with open(output_path, 'w') as f:
        for item in data:
            f.write(json.dumps(item) + '\n')
    
    print(f"Created {len(data)} samples at {output_path}")


if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", default="squ11z1/gpt-oss-9b-reasoning")
    parser.add_argument("--data", default=None)
    parser.add_argument("--output", default="./q_gpt_trained")
    parser.add_argument("--epochs", type=int, default=3)
    parser.add_argument("--create-data", action="store_true")
    
    args = parser.parse_args()
    
    if args.create_data:
        create_synthetic_training_data(args.model, args.data or "train_data.jsonl")
    else:
        train_quantum_head(
            model_name=args.model,
            train_data_path=args.data,
            output_dir=args.output,
            epochs=args.epochs,
        )