Blame for src/models.py

"""
Q-TensorFormer: Complete Model Architecture.

Unified Information-to-Resource Allocation Architecture.
Integrates:
  - Nested TT-FFN (ranks 1, 2, 4, 8 with instantaneous slicing)
  - Token Information State (8D feature tracking)
  - Information-Value Resource Allocator (marginal utility optimization)
  - Hardware-Aware Cost Profiling
  - Adaptive KV Cache with multi-precision & eviction
  - Deployment Presets:
      QTF_FULL, QTF_BALANCED, QTF_LATENCY, QTF_MEMORY,
      QTF_ENERGY, QTF_EDGE, QTF_CLASSICAL_ONLY
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import math
from typing import Optional, Dict, List, Tuple, Union

from .blocks import HybridBlock
from .config import ModelConfig, validate_model_config
from .kv_cache import AdaptiveKVCache
from .resource_allocator import AllocationBudget
from .hardware_cost_model import HardwareCostModel


class PositionalEncoding(nn.Module):
    def __init__(self, d_model: int, max_len: int = 512, dropout: float = 0.1):
        super().__init__()
        self.dropout = nn.Dropout(dropout)
        self.d_model = d_model
        pe = torch.zeros(max_len, d_model)
        pos = torch.arange(0, max_len).float().unsqueeze(1)
        div = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(pos * div)
        pe[:, 1::2] = torch.cos(pos * div)
        self.register_buffer("pe", pe.unsqueeze(0))

    def forward(self, x: torch.Tensor, start_pos: int = 0) -> torch.Tensor:
        seq_len = x.size(1)
        end_pos = start_pos + seq_len
        if end_pos <= self.pe.size(1):
            return self.dropout(x + self.pe[:, start_pos:end_pos, :])
        pos = torch.arange(start_pos, end_pos, device=x.device, dtype=torch.float32).unsqueeze(1)
        div = torch.exp(
            torch.arange(0, self.d_model, 2, device=x.device, dtype=torch.float32)
            * (-math.log(10000.0) / self.d_model)
        )
        pe_dyn = torch.zeros(1, seq_len, self.d_model, device=x.device, dtype=x.dtype)
        pe_dyn[0, :, 0::2] = torch.sin(pos * div)
        pe_dyn[0, :, 1::2] = torch.cos(pos * div)
        return self.dropout(x + pe_dyn)


class QTensorFormer(nn.Module):
    """
    Q-TensorFormer: Closed-Loop Information-to-Resource Adaptive Transformer.
    """

    DEPLOYMENT_PRESETS = [
        "QTF_FULL",
        "QTF_BALANCED",
        "QTF_LATENCY",
        "QTF_MEMORY",
        "QTF_ENERGY",
        "QTF_EDGE",
        "QTF_CLASSICAL_ONLY",
    ]

    def __init__(self, config: ModelConfig, preset: str = "QTF_BALANCED"):
        super().__init__()
        validate_model_config(config, self.__class__)
        self.config = config
        self.current_preset = preset.upper()

        self.embedding = nn.Embedding(config.vocab_size, config.d_model)
        self.pos_encoding = PositionalEncoding(config.d_model, config.max_seq_len, config.dropout)

        # Transformer blocks
        backend_type = getattr(config, "backend_type", "classical_surrogate")
        self.blocks = nn.ModuleList([
            HybridBlock(
                d_model=config.d_model,
                n_heads=config.n_heads,
                n_kv_heads=getattr(config, "n_kv_heads", None),
                ff_multiplier=config.ff_multiplier,
                tt_rank=config.tt_rank,
                tt_min_rank=config.tt_min_rank,
                use_quantum=config.use_quantum,
                n_qubits=config.n_qubits,
                backend_type=backend_type,
                dropout=config.dropout,
                max_seq_len=config.max_seq_len,
                hysteresis_tau=getattr(config, "hysteresis_tau", 0.15),
                default_preset=self.current_preset.replace("QTF_", "").lower(),
                enable_early_exit=getattr(config, "enable_early_exit", False),
                early_exit_threshold=getattr(config, "early_exit_threshold", 0.20),
                vocab_size=config.vocab_size,
            )
            for _ in range(config.n_layers)
        ])

        self.ln_f = nn.LayerNorm(config.d_model)
        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
        self.lm_head.weight = self.embedding.weight  # Weight tying

        # Hardware cost model
        self.hardware_model = HardwareCostModel()

        self._init_weights()

    def _init_weights(self):
        for name, p in self.named_parameters():
            if "weight" in name and p.dim() >= 2:
                nn.init.xavier_uniform_(p)
            elif "bias" in name:
                nn.init.zeros_(p)

    def set_preset(self, preset_name: str):
        """Switch deployment preset."""
        preset = preset_name.upper()
        if preset not in self.DEPLOYMENT_PRESETS:
            raise ValueError(f"Unknown preset {preset}. Choose from {self.DEPLOYMENT_PRESETS}")
        self.current_preset = preset

    def forward(
        self,
        input_ids: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        kv_caches: Optional[List[AdaptiveKVCache]] = None,
        budget: Optional[AllocationBudget] = None,
        return_stats: bool = False,
        preset: Optional[str] = None,
    ) -> Union[torch.Tensor, Tuple[torch.Tensor, List[Dict]]]:
        """
        Forward pass with dynamic information-to-resource allocation.
        """
        if preset is not None:
            self.set_preset(preset)

        B, T = input_ids.shape
        start_pos = kv_caches[0].seq_len if (kv_caches is not None and kv_caches[0].seq_len > 0) else 0

        x = self.embedding(input_ids)
        x = self.pos_encoding(x, start_pos=start_pos)

        preset_str = self.current_preset.replace("QTF_", "").lower()
        force_classical = (self.current_preset in ["QTF_CLASSICAL_ONLY", "QTF_EDGE"] or not self.config.use_quantum)

        all_stats = []
        for i, block in enumerate(self.blocks):
            layer_kv = kv_caches[i] if kv_caches is not None else None
            x, stats = block(
                x,
                mask=attention_mask,
                kv_cache=layer_kv,
                budget=budget,
                preset=preset_str,
                force_classical=force_classical,
            )
            all_stats.append(stats)
            if stats.get("early_exit_triggered", False) and "early_exit_logits" in stats:
                logits = stats["early_exit_logits"]
                if return_stats:
                    return logits, all_stats
                return logits

        x = self.ln_f(x)
        logits = self.lm_head(x)

        if return_stats:
            return logits, all_stats
        return logits

    @torch.no_grad()
    def generate(
        self,
        input_ids: torch.Tensor,
        max_new_tokens: int = 20,
        temperature: float = 1.0,
        top_k: int = 50,
        use_cache: bool = True,
        do_sample: bool = True,
        eos_token_id: Optional[int] = None,
        pad_token_id: Optional[int] = 0,
        attention_mask: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        """
        Autoregressive generation with KV Cache, context-window boundary handling,
        deterministic greedy mode (`do_sample=False` or `temperature=0.0`), and EOS/PAD stopping.
        """
        if input_ids is None or input_ids.dim() != 2 or input_ids.size(0) == 0 or input_ids.size(1) == 0:
            raise ValueError("input_ids must be a non-empty 2D tensor of shape (batch_size, seq_len).")
        if (input_ids < 0).any() or (input_ids >= self.config.vocab_size).any():
            raise ValueError(
                f"input_ids contains token IDs outside valid vocabulary range [0, {self.config.vocab_size - 1}]."
            )
        if max_new_tokens < 0:
            raise ValueError(f"max_new_tokens must be >= 0, got {max_new_tokens}.")
        if max_new_tokens == 0:
            return input_ids.clone()

        self.eval()
        B = input_ids.shape[0]
        max_ctx = max(1, int(self.config.max_seq_len))

        kv_caches = None
        if use_cache:
            kv_caches = [
                AdaptiveKVCache(max_capacity=max_ctx)
                for _ in range(self.config.n_layers)
            ]

        curr_input = input_ids
        curr_mask = attention_mask.clone() if attention_mask is not None else None
        finished = torch.zeros(B, dtype=torch.bool, device=input_ids.device)

        for _ in range(max_new_tokens):
            if kv_caches is not None and kv_caches[0].seq_len > 0:
                step_input = curr_input[:, -1:]
                # When KV cache is active and may evict to max_ctx, align mask to cached seq_len + 1
                if curr_mask is not None:
                    expected_k_len = min(max_ctx, kv_caches[0].seq_len + 1)
                    step_mask = curr_mask[:, -expected_k_len:]
                else:
                    step_mask = None
            else:
                # Prefill or uncached step: enforce context-window boundary max_seq_len
                step_input = curr_input[:, -max_ctx:]
                step_mask = curr_mask[:, -max_ctx:] if curr_mask is not None else None

            logits = self(step_input, attention_mask=step_mask, kv_caches=kv_caches)
            next_logits = logits[:, -1, :]

            if not do_sample or temperature <= 0.0:
                next_token = torch.argmax(next_logits, dim=-1, keepdim=True)
            else:
                scaled_logits = next_logits / max(1e-5, float(temperature))
                if top_k > 0:
                    v, _ = torch.topk(scaled_logits, min(top_k, scaled_logits.size(-1)))
                    scaled_logits = scaled_logits.masked_fill(scaled_logits < v[:, [-1]], float("-inf"))
                probs = F.softmax(scaled_logits, dim=-1)
                next_token = torch.multinomial(probs, num_samples=1)

            if eos_token_id is not None:
                if pad_token_id is not None:
                    next_token = torch.where(
                        finished.unsqueeze(-1),
                        torch.full_like(next_token, pad_token_id),
                        next_token,
                    )
                finished = finished | (next_token.squeeze(-1) == eos_token_id)

            curr_input = torch.cat([curr_input, next_token], dim=-1)
            if curr_mask is not None:
                new_mask_col = (~finished).long().unsqueeze(-1) if eos_token_id is not None else torch.ones(
                    (B, 1), dtype=curr_mask.dtype, device=curr_mask.device
                )
                curr_mask = torch.cat([curr_mask, new_mask_col.to(curr_mask.dtype)], dim=-1)

            if eos_token_id is not None and finished.all():
                break

        return curr_input

    def reset_schedulers(self):
        for block in self.blocks:
            block.reset_scheduler()

    @property
    def total_params(self) -> int:
        return sum(p.numel() for p in self.parameters())

    @property
    def active_params(self) -> int:
        tt_active = sum(b.tt_ffn.active_params for b in self.blocks)
        base = self.total_params - sum(b.tt_ffn.total_params for b in self.blocks)
        return base + tt_active

    @property
    def compression_ratio(self) -> float:
        dense_per_block = 2 * self.config.d_model * self.config.d_model * self.config.ff_multiplier
        base = self.total_params - sum(b.tt_ffn.total_params for b in self.blocks)
        active_tt = sum(b.tt_ffn.active_params for b in self.blocks)
        return (base + dense_per_block * self.config.n_layers) / max(base + active_tt, 1)

    def flops_estimate(self, batch_size: int = 1, seq_len: int = 32) -> Dict:
        total = 0
        breakdown = {}
        for i, block in enumerate(self.blocks):
            b = block.flops_estimate(batch_size, seq_len)
            total += b["total"]
            breakdown[f"block_{i}"] = b
        return {"total": total, "breakdown": breakdown}


class DenseBaseline(nn.Module):
    """
    Standard dense transformer baseline with identical hyperparameters.
    """

    def __init__(self, config: ModelConfig):
        super().__init__()
        validate_model_config(config, self.__class__)
        self.config = config

        self.embedding = nn.Embedding(config.vocab_size, config.d_model)
        self.pos_encoding = PositionalEncoding(config.d_model, config.max_seq_len, config.dropout)

        self.blocks = nn.ModuleList([
            nn.ModuleDict({
                "ln1": nn.LayerNorm(config.d_model),
                "attn": nn.MultiheadAttention(config.d_model, config.n_heads, dropout=config.dropout, batch_first=True),
                "ln2": nn.LayerNorm(config.d_model),
                "ffn": nn.Sequential(
                    nn.Linear(config.d_model, config.d_model * config.ff_multiplier),
                    nn.GELU(),
                    nn.Linear(config.d_model * config.ff_multiplier, config.d_model),
                ),
                "dropout": nn.Dropout(config.dropout),
            })
            for _ in range(config.n_layers)
        ])

        self.ln_f = nn.LayerNorm(config.d_model)
        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
        self.lm_head.weight = self.embedding.weight
        self._init_weights()

    def _init_weights(self):
        for name, p in self.named_parameters():
            if "weight" in name and p.dim() >= 2:
                nn.init.xavier_uniform_(p)
            elif "bias" in name:
                nn.init.zeros_(p)

    def forward(self, input_ids, attention_mask=None, return_stats=False):
        B, T = input_ids.shape
        x = self.embedding(input_ids)
        x = self.pos_encoding(x)

        causal_hide = torch.triu(torch.ones(T, T, device=x.device, dtype=torch.bool), diagonal=1)
        if attention_mask is None:
            attn_mask = torch.zeros((T, T), device=x.device, dtype=x.dtype).masked_fill(causal_hide, float("-inf"))
        else:
            m = attention_mask.to(device=x.device)
            pad_hide = (~m) if m.dtype == torch.bool else (m == 0)
            combined_hide = causal_hide.unsqueeze(0) | pad_hide[:, None, :T]
            row_all_masked = combined_hide.all(dim=-1, keepdim=True)
            diag = torch.eye(T, device=x.device, dtype=torch.bool).unsqueeze(0)
            safe_hide = torch.where(row_all_masked, ~diag, combined_hide)
            attn_mask = torch.zeros((B, T, T), device=x.device, dtype=x.dtype).masked_fill(safe_hide, float("-inf"))
            attn_mask = attn_mask.repeat_interleave(self.config.n_heads, dim=0)

        for block in self.blocks:
            normed = block["ln1"](x)
            attn_out, _ = block["attn"](normed, normed, normed, attn_mask=attn_mask, need_weights=False)
            x = x + block["dropout"](attn_out)
            ffn_out = block["ffn"](block["ln2"](x))
            x = x + block["dropout"](ffn_out)

        x = self.ln_f(x)
        logits = self.lm_head(x)
        if return_stats:
            return logits, []
        return logits

    @property
    def total_params(self) -> int:
        return sum(p.numel() for p in self.parameters())


def create_model(config: ModelConfig, model_type: str = "qtensor", preset: str = "QTF_BALANCED"):
    """Factory helper to build models."""
    if model_type in ["qtensor", "hybrid"]:
        return QTensorFormer(config, preset=preset)
    elif model_type in ["dense", "baseline"]:
        return DenseBaseline(config)
    elif model_type == "tensor_only":
        cfg = ModelConfig(**config.__dict__)
        cfg.use_quantum = False
        return QTensorFormer(cfg, preset="QTF_CLASSICAL_ONLY")
    else:
        raise ValueError(f"Unknown model_type: {model_type}")