"""Standalone TaoNet model used by the HF wrapper.""" from types import SimpleNamespace import torch import torch.nn as nn import torch.nn.functional as F try: from .embeddings import FactorizedEmbedding from .mla_components import AttentionBlock except ImportError: from embeddings import FactorizedEmbedding from mla_components import AttentionBlock class SimpleLLM(nn.Module): """Pure attention TaoNet language model.""" def __init__(self, config): super().__init__() self.config = config self.vocab_size = config.vocab_size self.d_model = config.hidden_dim self.n_layers = config.num_layers self.n_heads = config.num_heads self.dropout = config.dropout self.d_latent_kv = config.d_latent_kv if config.d_latent_kv is not None else int(self.d_model * 0.75) self.d_rope = config.d_rope if config.d_rope is not None else (self.d_model // self.n_heads) self.d_ff = config.hidden_dim_ff if config.hidden_dim_ff is not None else (self.d_model * 4) self.gqa_groups = getattr(config, "gqa_groups", 1) self.use_factorized_embedding = getattr(config, "use_factorized_embedding", False) self.d_embed_rank = getattr(config, "d_embed_rank", 96) self.rope_scale = getattr(config, "rope_scale", 40.0) self.yarn_enabled = getattr(config, "yarn_enabled", False) self.yarn_original_max_seq_length = getattr(config, "yarn_original_max_seq_length", None) self.yarn_alpha = getattr(config, "yarn_alpha", 1.0) self.max_seq_length = config.max_seq_length if self.d_model % self.n_heads != 0: raise ValueError("hidden_dim must be divisible by num_heads") if self.d_latent_kv % self.n_heads != 0: raise ValueError("d_latent_kv must be divisible by num_heads") if self.use_factorized_embedding: self.token_embedding = FactorizedEmbedding( self.vocab_size, self.d_model, self.d_embed_rank, ) else: self.token_embedding = nn.Embedding(self.vocab_size, self.d_model) self.embedding_dropout = nn.Dropout(self.dropout) self.blocks = nn.ModuleList( [ AttentionBlock( d_model=self.d_model, d_latent_kv=self.d_latent_kv, n_heads=self.n_heads, d_rope=self.d_rope, d_ff=int(self.d_ff), dropout=self.dropout, gqa_groups=self.gqa_groups, rope_scale=self.rope_scale, max_seq_length=self.max_seq_length, yarn_enabled=self.yarn_enabled, yarn_original_max_seq_length=self.yarn_original_max_seq_length, yarn_alpha=self.yarn_alpha, ) for _ in range(self.n_layers) ] ) self.final_norm = nn.LayerNorm(self.d_model) self.output_head = nn.Linear(self.d_model, self.vocab_size, bias=False) self.register_buffer("causal_mask_cache", None, persistent=False) self.apply(self._init_weights) def _init_weights(self, module): if isinstance(module, nn.Linear): nn.init.normal_(module.weight, mean=0.0, std=0.02) if module.bias is not None: nn.init.zeros_(module.bias) elif isinstance(module, nn.Embedding): nn.init.normal_(module.weight, mean=0.0, std=0.02) def _get_causal_mask(self, seq_len, device): if self.causal_mask_cache is None or self.causal_mask_cache.size(-1) < seq_len: mask = torch.tril(torch.ones(seq_len, seq_len, device=device, dtype=torch.bool)) self.register_buffer("causal_mask_cache", mask, persistent=False) return self.causal_mask_cache[:seq_len, :seq_len] def forward( self, input_ids=None, attention_mask=None, labels=None, inputs_embeds=None, pixel_values=None, ): del pixel_values if inputs_embeds is None: if input_ids is None: raise ValueError("Either input_ids or inputs_embeds must be provided") _, seq_len = input_ids.shape device = input_ids.device x = self.token_embedding(input_ids) else: _, seq_len, _ = inputs_embeds.shape device = inputs_embeds.device x = inputs_embeds causal_mask = self._get_causal_mask(seq_len, device) if attention_mask is not None: padding_mask = attention_mask.unsqueeze(1).unsqueeze(1).bool() combined_mask = causal_mask.unsqueeze(0).unsqueeze(0) & padding_mask combined_mask = combined_mask.float() else: combined_mask = causal_mask.unsqueeze(0).unsqueeze(0).float() x = self.embedding_dropout(x) for block in self.blocks: x = block(x, attention_mask=combined_mask) x = self.final_norm(x) logits = self.output_head(x) loss = None if labels is not None: logits_flat = logits.view(-1, logits.size(-1)) labels_flat = labels.view(-1) if not torch.any(labels_flat != -100): raise ValueError("All labels are masked out (-100), so loss cannot be computed") loss = F.cross_entropy(logits_flat, labels_flat, reduction="mean", ignore_index=-100) return {"logits": logits, "loss": loss} def build_runtime_config(config): """Create a light-weight runtime config object for the standalone model.""" return SimpleNamespace( vocab_size=config.vocab_size, hidden_dim=config.hidden_dim, num_layers=config.num_layers, num_heads=config.num_heads, dropout=config.dropout, max_seq_length=config.max_seq_length, d_latent_kv=config.d_latent_kv, d_rope=config.d_rope, gqa_groups=config.gqa_groups, hidden_dim_ff=config.hidden_dim_ff, use_factorized_embedding=config.use_factorized_embedding, d_embed_rank=config.d_embed_rank, rope_scale=config.rope_scale, yarn_enabled=config.yarn_enabled, yarn_original_max_seq_length=config.yarn_original_max_seq_length, yarn_alpha=config.yarn_alpha, )