cat_v3/model.py

"""Integrated CAT V3 model architecture."""

from __future__ import annotations

from typing import Any, Dict, List, Optional
import torch
import torch.nn as nn

from cat_v3.encoder import TinyEncoder
from cat_v3.router import SemanticRouter
from cat_v3.experts import GATExpert
from cat_v3.fusion import ConceptFusionLayer
from cat_v3.combiner import TinyCombiner
from cat_v3.decoder import TinyDecoder
from cat_v3.dataset import DOMAINS


class ConceptMemory(nn.Module):
    """Shared concept memory embeddings table."""

    def __init__(self, num_concepts: int, concept_dim: int) -> None:
        super().__init__()
        self.embeddings = nn.Embedding(num_concepts, concept_dim)

    def forward(self) -> torch.Tensor:
        """Returns all concept embeddings."""
        ids = torch.arange(self.embeddings.num_embeddings, device=self.embeddings.weight.device)
        return self.embeddings(ids)


class CATV3Model(nn.Module):
    """Concept Attention Transformer V3 (CAT V3) integrated model."""

    def __init__(
        self,
        num_concepts: int,
        tokenizer_vocab_size: int,
        pad_id: int,
        eos_id: int,
        expert_graphs: Dict[str, Tuple[torch.Tensor, torch.Tensor]],
        concept_dim: int = 128,
        hidden_size: int = 128,
        path_length: int = 8,
        top_m: int = 8,
        decoder_vocab_size: int = 500,
    ) -> None:
        super().__init__()
        self.num_concepts = num_concepts
        self.concept_dim = concept_dim
        self.pad_id = pad_id
        self.eos_id = eos_id
        self.path_length = path_length
        self.top_m = top_m

        # 1. Tiny Encoder
        self.encoder = TinyEncoder(vocab_size=tokenizer_vocab_size, hidden_size=hidden_size)

        # 2. Semantic Router
        self.router = SemanticRouter(encoder_dim=hidden_size, num_experts=len(DOMAINS))

        # 3. Concept Memory
        self.memory = ConceptMemory(num_concepts=num_concepts, concept_dim=concept_dim)
        
        # Project encoder representation into concept space if dimensions differ
        self.query_proj = nn.Linear(hidden_size, concept_dim)

        # 4. GAT Experts registered as a ModuleDict
        self.experts = nn.ModuleDict()
        for domain in DOMAINS:
            edge_index, edge_weight = expert_graphs[domain]
            self.experts[domain] = GATExpert(
                domain_name=domain,
                num_concepts=num_concepts,
                concept_dim=concept_dim,
                edge_index=edge_index,
                edge_weight=edge_weight,
                pad_id=pad_id,
                eos_id=eos_id,
                path_length=path_length,
            )

        # 5. Concept Fusion Layer
        self.fusion = ConceptFusionLayer(num_concepts=num_concepts, pad_id=pad_id, eos_id=eos_id, top_m=top_m)

        # 6. Tiny Combiner
        self.combiner = TinyCombiner(concept_dim=concept_dim)

        # 7. Tiny Decoder
        self.decoder = TinyDecoder(
            vocab_size=decoder_vocab_size,
            concept_dim=concept_dim,
            hidden_size=hidden_size,
        )

    def forward(
        self,
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor,
        target_paths: Optional[torch.Tensor] = None,
        target_responses: Optional[torch.Tensor] = None,
        target_responses_mask: Optional[torch.Tensor] = None,
        router_top_k: int = 2,
        router_threshold: float = 0.5,
    ) -> Dict[str, Any]:
        """Runs the complete forward pass of CAT V3.
        
        Args:
            input_ids: [batch_size, seq_len]
            attention_mask: [batch_size, seq_len]
            target_paths: [batch_size, path_length] (optional)
            target_responses: [batch_size, resp_len] (optional)
            target_responses_mask: [batch_size, resp_len] (optional)
            router_top_k: Minimum number of experts to activate
            router_threshold: Activation threshold for experts
            
        Returns:
            Dict containing outputs of all sub-components.
        """
        # Step 1: Tiny Encoder
        query_emb = self.encoder(input_ids, attention_mask)
        
        # Step 2: Semantic Router
        router_out = self.router(query_emb, top_k=router_top_k, threshold=router_threshold)
        router_logits = router_out["logits"]
        router_probs = router_out["probs"]
        router_mask = router_out["activation_mask"]

        # Step 3: Global Concept Embeddings
        global_embs = self.memory()
        query_context = self.query_proj(query_emb)

        # Step 4: Graph Mixture of Experts (Graph-MoE)
        expert_reports = {}
        for domain in DOMAINS:
            expert = self.experts[domain]
            # Pass through expert
            expert_reports[domain] = expert(
                global_embeddings=global_embs,
                query_context=query_context,
                target_paths=target_paths,
            )

        # Step 5: Concept Fusion Layer
        fusion_out = self.fusion(
            expert_reports=expert_reports,
            router_probs=router_probs,
            router_mask=router_mask,
            global_embeddings=global_embs,
            domain_names=DOMAINS,
        )
        fused_embs = fusion_out["fused_embeddings"]

        # Step 6: Tiny Combiner
        organized_embs = self.combiner(fused_embs)

        # Step 7: Tiny Decoder
        decoder_logits = None
        if target_responses is not None:
            # We predict next tokens, so shift targets during training
            decoder_logits = self.decoder(
                organized_embeddings=organized_embs,
                target_ids=target_responses,
                target_mask=target_responses_mask,
            )

        return {
            "query_embedding": query_emb,
            "router_logits": router_logits,
            "router_probs": router_probs,
            "router_mask": router_mask,
            "expert_reports": expert_reports,
            "fused_concept_ids": fusion_out["fused_concept_ids"],
            "fused_scores": fusion_out["fused_scores"],
            "fused_embeddings": fused_embs,
            "organized_embeddings": organized_embs,
            "decoder_logits": decoder_logits,
        }

    @torch.no_grad()
    def generate_response(
        self,
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor,
        max_length: int = 32,
        router_top_k: int = 2,
        router_threshold: float = 0.5,
    ) -> Dict[str, Any]:
        """Runs inference to generate a reasoning graph and natural language answer."""
        # 1. Forward pass to get organized embeddings
        outputs = self.forward(
            input_ids=input_ids,
            attention_mask=attention_mask,
            router_top_k=router_top_k,
            router_threshold=router_threshold,
        )
        
        organized_embs = outputs["organized_embeddings"]
        router_mask = outputs["router_mask"]
        expert_reports = outputs["expert_reports"]

        # 2. Autoregressive token generation
        gen_tokens = self.decoder.generate(
            organized_embeddings=organized_embs,
            max_length=max_length,
            start_id=self.pad_id,
            eos_id=self.eos_id,
        )

        return {
            "router_probs": outputs["router_probs"],
            "router_mask": router_mask,
            "expert_reports": expert_reports,
            "fused_concept_ids": outputs["fused_concept_ids"],
            "generated_tokens": gen_tokens,
        }