training/models.py

"""
Model components for Speech-to-Speech training.

Single Responsibility: Only defines model architectures.
Open/Closed: Can extend with new adapters without modifying existing code.

Optimizations:
- Flash Attention 2 for memory-efficient attention (10-20x savings on long sequences)
- BFloat16 for better numerical stability than FP16
- Gradient checkpointing for memory savings
"""

import torch
import torch.nn as nn
from typing import Optional
from .config import (
    DEFAULT_WHISPER_DIM,
    DEFAULT_LLM_DIM,
    DEFAULT_DOWNSAMPLE,
    DEFAULT_INTERMEDIATE_DIM,
)
from .utils import log


class SpeechAdapter(nn.Module):
    """
    Speech adapter that maps Whisper features to LLM embedding space.

    Architecture: 5× downsampling + FFN with intermediate dim

    Based on LLaMA-Omni 2 design:
    - Concatenates 5 consecutive Whisper frames
    - Projects through 2-layer FFN
    - Applies LayerNorm for stability

    Args:
        whisper_dim: Dimension of Whisper features (default: 1280)
        llm_dim: Dimension of LLM embeddings (default: 3072)
        downsample: Downsampling factor (default: 5)
        intermediate_dim: Hidden dimension of FFN (default: 2048)
    """

    def __init__(
        self,
        whisper_dim: int = DEFAULT_WHISPER_DIM,
        llm_dim: int = DEFAULT_LLM_DIM,
        downsample: int = DEFAULT_DOWNSAMPLE,
        intermediate_dim: int = DEFAULT_INTERMEDIATE_DIM
    ):
        super().__init__()
        self.whisper_dim = whisper_dim
        self.llm_dim = llm_dim
        self.downsample = downsample
        self.intermediate_dim = intermediate_dim

        concat_dim = whisper_dim * downsample

        self.ffn = nn.Sequential(
            nn.Linear(concat_dim, intermediate_dim),
            nn.GELU(),
            nn.Linear(intermediate_dim, llm_dim),
            nn.LayerNorm(llm_dim)
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass.

        Args:
            x: Whisper features [B, T, D]

        Returns:
            LLM embeddings [B, T // downsample, llm_dim]
        """
        B, T, D = x.shape

        # Ensure T is divisible by downsample
        T_new = (T // self.downsample) * self.downsample
        x = x[:, :T_new]

        # Reshape: [B, T, D] -> [B, T // downsample, D * downsample]
        x = x.reshape(B, T_new // self.downsample, D * self.downsample)

        return self.ffn(x)

    def get_num_params(self) -> int:
        """Return total number of parameters."""
        return sum(p.numel() for p in self.parameters())

    def get_config(self) -> dict:
        """Return configuration dict for serialization."""
        return {
            "whisper_dim": self.whisper_dim,
            "llm_dim": self.llm_dim,
            "downsample": self.downsample,
            "intermediate_dim": self.intermediate_dim,
        }

    @classmethod
    def from_config(cls, config: dict) -> 'SpeechAdapter':
        """Create adapter from configuration dict."""
        return cls(**config)


class ModelFactory:
    """
    Factory for creating models with consistent settings.

    Single Responsibility: Only handles model instantiation.
    Dependency Inversion: Depends on abstractions (config), not concretions.
    """

    @staticmethod
    def create_adapter(
        whisper_dim: int = DEFAULT_WHISPER_DIM,
        llm_dim: int = DEFAULT_LLM_DIM,
        dtype: torch.dtype = torch.float32,
        checkpoint_path: Optional[str] = None
    ) -> SpeechAdapter:
        """
        Create a SpeechAdapter, optionally loading from checkpoint.

        Args:
            whisper_dim: Whisper feature dimension
            llm_dim: LLM embedding dimension
            dtype: Tensor dtype
            checkpoint_path: Optional path to checkpoint

        Returns:
            Initialized SpeechAdapter
        """
        adapter = SpeechAdapter(
            whisper_dim=whisper_dim,
            llm_dim=llm_dim,
        ).to(dtype=dtype)

        if checkpoint_path:
            ckpt = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
            if "adapter" in ckpt:
                adapter.load_state_dict(ckpt["adapter"])
            elif "state_dict" in ckpt:
                adapter.load_state_dict(ckpt["state_dict"])
            else:
                adapter.load_state_dict(ckpt)

        return adapter

    @staticmethod
    def create_llm(
        model_path: str,
        dtype: torch.dtype = torch.bfloat16,
        freeze: bool = True,
        gradient_checkpointing: bool = False,
        use_flash_attention: bool = True,
        verbose: bool = True,
    ):
        """
        Create and configure the LLM with memory optimizations.

        Args:
            model_path: HuggingFace model path
            dtype: Tensor dtype (BF16 recommended for stability)
            freeze: Whether to freeze all parameters
            gradient_checkpointing: Enable gradient checkpointing
            use_flash_attention: Try to use Flash Attention 2 (10-20x memory savings)
            verbose: Log configuration details

        Returns:
            Configured LLM model
        """
        from transformers import AutoModelForCausalLM

        # Determine attention implementation
        # Flash Attention 2 provides 10-20x memory savings on long sequences
        # Requires: Ampere/Ada/Hopper GPU (RTX 30xx, 40xx, A100, H100)
        attn_impl = "sdpa"  # Default fallback

        if use_flash_attention and torch.cuda.is_available():
            try:
                # Actually try to import flash_attn to verify it works
                import flash_attn
                # Check GPU capability (Flash Attention 2 requires SM 80+)
                major, _ = torch.cuda.get_device_capability()
                if major >= 8:  # Ampere or newer (RTX 30xx, 40xx, A100, H100)
                    attn_impl = "flash_attention_2"
                    if verbose:
                        log(f"[LLM] Using Flash Attention 2 v{flash_attn.__version__} (10-20x memory savings)")
                else:
                    if verbose:
                        log(f"[LLM] GPU SM {major}.x too old for Flash Attention 2, using SDPA")
            except ImportError:
                if verbose:
                    log("[LLM] flash_attn not installed, using SDPA")
            except Exception as e:
                if verbose:
                    log(f"[LLM] Flash Attention check failed: {e}, using SDPA")

        # Load model with optimizations
        try:
            llm = AutoModelForCausalLM.from_pretrained(
                model_path,
                torch_dtype=dtype,
                attn_implementation=attn_impl,
            )
        except Exception as e:
            # Fallback if flash_attention_2 fails
            if attn_impl == "flash_attention_2":
                if verbose:
                    log(f"[LLM] Flash Attention failed ({e}), falling back to SDPA")
                llm = AutoModelForCausalLM.from_pretrained(
                    model_path,
                    torch_dtype=dtype,
                    attn_implementation="sdpa",
                )
            else:
                raise

        if freeze:
            for p in llm.parameters():
                p.requires_grad = False
            llm.eval()

        if gradient_checkpointing:
            llm.gradient_checkpointing_enable()
            if verbose:
                log("[LLM] Gradient checkpointing enabled")

        return llm

    @staticmethod
    def apply_lora(
        llm,
        lora_config: 'LoRAConfig'
    ):
        """
        Apply LoRA to an LLM.

        Args:
            llm: The LLM model
            lora_config: LoRA configuration

        Returns:
            LLM with LoRA applied
        """
        from peft import get_peft_model
        peft_config = lora_config.to_peft_config()
        return get_peft_model(llm, peft_config)