Upload cat_v3/model.py with huggingface_hub

3b538c8 verified 5 months ago

7.54 kB

"""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,
        }

	"""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,
	}