cat_v3/train.py

"""Training pipeline for CAT V3."""

from __future__ import annotations

import os
from typing import Dict, List, Optional
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader

from cat_v3.dataset import (
    DOMAINS,
    ConceptVocabulary,
    SimpleCharTokenizer,
    CATV3Dataset,
    grow_dataset,
    build_expert_graphs
)
from cat_v3.model import CATV3Model


def cat_v3_loss(
    outputs: Dict[str, torch.Tensor],
    router_target: torch.Tensor,
    path_ids: torch.Tensor,
    response_targets: torch.Tensor,
    pad_id: int,
    model: nn.Module,
) -> Tuple[torch.Tensor, Dict[str, float]]:
    """Computes joint multi-task loss: Router BCE + GAT Experts CE + Causal Decoder CE."""
    batch_size = router_target.size(0)
    num_concepts = outputs["fused_scores"].size(1)
    
    # 1. Router Loss (BCE)
    router_logits = outputs["router_logits"]
    loss_router = F.binary_cross_entropy_with_logits(router_logits, router_target)
    
    # 2. Experts Loss (Cross-Entropy on paths for active experts)
    loss_experts = torch.tensor(0.0, device=router_target.device)
    active_expert_count = 0
    
    for idx, domain in enumerate(DOMAINS):
        expert_out = outputs["expert_reports"][domain]
        expert_logits = expert_out["path_logits"]  # [B, path_length, num_concepts]
        path_len = expert_logits.size(1)
        
        # Extract filtered target path for this expert
        expert_target_path = expert_out.get("target_path")
        if expert_target_path is None:
            expert_target_path = path_ids[:, :path_len]
            
        # Calculate CE step-by-step against the filtered domain-specific path
        ce = F.cross_entropy(
            expert_logits.view(-1, num_concepts),
            expert_target_path.reshape(-1),
            reduction="none"
        ).view(batch_size, path_len).mean(dim=1)  # [B]
        
        # Mask out samples where this expert has no concepts in the path
        expert_module = model.experts[domain]
        domain_nodes = set(expert_module.edge_index.view(-1).cpu().tolist())
        domain_nodes.discard(expert_module.eos_id)
        domain_nodes.discard(expert_module.pad_id)
        
        has_domain_concept = torch.zeros(batch_size, dtype=torch.float, device=router_target.device)
        for b in range(batch_size):
            for val in expert_target_path[b].tolist():
                if val in domain_nodes:
                    has_domain_concept[b] = 1.0
                    break
                    
        # Active weight is router activation * has_domain_concept
        active_weight = router_target[:, idx] * has_domain_concept
        if active_weight.sum() > 0:
            loss_experts = loss_experts + (ce * active_weight).sum() / active_weight.sum()
            active_expert_count += 1
            
    if active_expert_count > 0:
        loss_experts = loss_experts / active_expert_count
        
    # 3. Decoder Loss (Cross-Entropy on target words)
    decoder_logits = outputs["decoder_logits"]  # [B, seq_len-1, vocab_size]
    loss_decoder = F.cross_entropy(
        decoder_logits.reshape(-1, decoder_logits.size(-1)),
        response_targets.reshape(-1),
        ignore_index=pad_id
    )
    
    # Total loss
    total_loss = 0.5 * loss_router + 1.0 * loss_experts + 1.0 * loss_decoder
    
    return total_loss, {
        "loss_total": total_loss.item(),
        "loss_router": loss_router.item(),
        "loss_experts": loss_experts.item(),
        "loss_decoder": loss_decoder.item()
    }


def train_cat_v3(
    epochs: int = 20,
    batch_size: int = 8,
    lr: float = 5e-4,
    checkpoint_dir: str = "checkpoints/cat_v3",
) -> CATV3Model:
    """Trains the full CAT V3 prototype model."""
    os.makedirs(checkpoint_dir, exist_ok=True)
    
    # 1. Prepare data
    raw_data = grow_dataset()
    
    # Get all texts to fit the decoder tokenizer vocab
    all_texts = [d["question"] for d in raw_data] + [d["response"] for d in raw_data]
    tokenizer = SimpleCharTokenizer(all_texts)
    
    vocab = ConceptVocabulary({d: list(set(grow_dataset()[0]["concept_paths"][0])) for d in DOMAINS})
    # Re-extract all concepts from dataset to build clean vocabulary
    concept_lists = {}
    for d in DOMAINS:
        concept_lists[d] = []
    
    for item in raw_data:
        # Sort concepts by domain
        for domain in item["active_experts"]:
            for path in item["concept_paths"]:
                for concept in path:
                    concept_lists[domain].append(concept)
                    
    # Ensure every domain has some concepts
    for d in DOMAINS:
        concept_lists[d] = list(set(concept_lists[d]))
        if not concept_lists[d]:
            # fallback
            concept_lists[d] = ["entropy", "load"]
            
    vocab = ConceptVocabulary(concept_lists)
    expert_graphs = build_expert_graphs(vocab, raw_data)
    
    dataset = CATV3Dataset(raw_data, concept_vocab=vocab, token_tokenizer=tokenizer)
    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
    
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Training CAT V3 on device: {device}")
    
    # 2. Initialize Model
    model = CATV3Model(
        num_concepts=vocab.size(),
        tokenizer_vocab_size=tokenizer.vocab_size(),
        pad_id=tokenizer.pad_id,
        eos_id=tokenizer.eos_id,
        expert_graphs=expert_graphs,
        concept_dim=128,
        hidden_size=128,
        path_length=8,
        top_m=8,
        decoder_vocab_size=tokenizer.vocab_size()
    ).to(device)
    
    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-4)
    
    # 3. Training loop
    model.train()
    for epoch in range(epochs):
        epoch_losses = {"loss_total": 0.0, "loss_router": 0.0, "loss_experts": 0.0, "loss_decoder": 0.0}
        steps = 0
        
        for batch in dataloader:
            input_ids = batch["input_ids"].to(device)
            attention_mask = batch["attention_mask"].to(device)
            router_target = batch["router_target"].to(device)
            path_ids = batch["path_ids"].to(device)
            response_ids = batch["response_ids"].to(device)
            response_mask = batch["response_mask"].to(device)
            
            # Shift targets for causal language modeling
            decoder_inputs = response_ids[:, :-1]
            decoder_targets = response_ids[:, 1:]
            decoder_attention_mask = response_mask[:, :-1]
            
            optimizer.zero_grad()
            
            outputs = model(
                input_ids=input_ids,
                attention_mask=attention_mask,
                target_paths=path_ids,
                target_responses=decoder_inputs,
                target_responses_mask=decoder_attention_mask,
            )
            
            loss, metrics = cat_v3_loss(
                outputs=outputs,
                router_target=router_target,
                path_ids=path_ids,
                response_targets=decoder_targets,
                pad_id=tokenizer.pad_id,
                model=model
            )
            
            loss.backward()
            optimizer.step()
            
            for k, v in metrics.items():
                epoch_losses[k] += v
            steps += 1
            
        # Log epoch summary
        if (epoch + 1) % 5 == 0 or epoch == 0:
            avg_loss = {k: v / steps for k, v in epoch_losses.items()}
            print(
                f"Epoch {epoch+1:02d}/{epochs:02d} - "
                f"Total: {avg_loss['loss_total']:.4f} | "
                f"Router: {avg_loss['loss_router']:.4f} | "
                f"Experts: {avg_loss['loss_experts']:.4f} | "
                f"Decoder: {avg_loss['loss_decoder']:.4f}"
            )
            
    # Save checkpoint
    checkpoint_path = os.path.join(checkpoint_dir, "cat_v3_model.pt")
    torch.save({
        "model_state_dict": model.state_dict(),
        "vocab": vocab,
        "tokenizer": tokenizer,
        "expert_graphs": expert_graphs,
    }, checkpoint_path)
    print(f"Model saved to {checkpoint_path}")
    
    return model