src/model.py

"""
Model Architecture: Lightweight Projection Head.

This module implements the trainable parameters of the system: a multi-layer
perceptron (MLP) with dropout and LayerNorm that maps change features to a
multimodal space. The CLIP backbone is frozen — only this lightweight adapter is
learned. The contrastive loss that trains it lives with each trainer
(``train._masked_infonce`` for the PEFT ProjectionHead; ``lora_train._infonce_loss``
for LoRA).
"""
from pathlib import Path
from typing import Any, Dict, Optional, Tuple

import torch
import torch.nn as nn


class ProjectionHead(nn.Module):
    """
    Lightweight MLP adapter that maps change features to multimodal space.

    This is the only trainable component of the system. The frozen CLIP vision/backbone
    encoders are passed through a small neural network with dropout for regularization
    and LayerNorm for training stability.

    Architecture:
        Input (delta_f): [batch_size, embed_dim]  e.g., [B, 768]
                          - For difference mode: same dimension as CLIP embeddings
                          - For concatenation: double the dimension (e.g., 1536)
        Hidden layers: ReLU + Dropout + LayerNorm (repeated for capacity)
        Output: [batch_size, embed_dim]  Projected to match text embedding space

    Example:
        >>> model = ProjectionHead(input_dim=768, hidden_dim=512, output_dim=768)
        >>> delta_f = torch.randn(32, 768)  # batch of 32 change pairs
        >>> projected = model(delta_f)
        >>> print(projected.shape)  # torch.Size([32, 768])
    """

    def __init__(
        self,
        input_dim: int,
        output_dim: Optional[int] = None,
        hidden_dims: Tuple[int, ...] = (512, 256),
        dropout_rate: float = 0.3,
        layer_norm_eps: float = 1e-6
    ):
        """
        Initialize the projection head.

        Args:
            input_dim (int): Dimensionality of change features (delta_f).
                Typically matches CLIP embedding dimension (768 for ViT-L/14).
            output_dim (int, optional): Output dimension. Defaults to input_dim.
            hidden_dims (tuple): Intermediate layer dimensions. More layers = more
time but better capacity.
            dropout_rate (float): Dropout probability after each ReLU. 0.3 is a good
                starting point for regularization without underfitting.
            layer_norm_eps: Epsilon for numerical stability in LayerNorm.
        """
        super().__init__()

        self.input_dim = input_dim
        self.output_dim = output_dim if output_dim is not None else input_dim
        self.hidden_dims = hidden_dims

        # Build the MLP with alternating ReLU, Dropout, LayerNorm layers
        layers = []
        current_dim = input_dim

        for i, hidden_dim in enumerate(hidden_dims):
            # Linear projection + ReLU activation
            layers.append(nn.Linear(current_dim, hidden_dim))
            layers.append(nn.ReLU())
            # LayerNorm after each non-linearity (stabilizes training)
            layers.append(nn.LayerNorm(hidden_dim, eps=layer_norm_eps))
            # Dropout for regularization
            if dropout_rate > 0:
                layers.append(nn.Dropout(dropout_rate))

            current_dim = hidden_dim

        # Final projection to output_dim with LayerNorm
        layers.append(nn.Linear(current_dim, self.output_dim))
        layers.append(nn.LayerNorm(self.output_dim, eps=layer_norm_eps))

        self.mlp = nn.Sequential(*layers)

        # Initialize weights with Xavier/Glorot initialization (better for ReLU)
        for module in self.modules():
            if isinstance(module, nn.Linear):
                nn.init.xavier_uniform_(module.weight)
                nn.init.zeros_(module.bias)

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

        Args:
            x (torch.Tensor): Input change features. Shape: [batch_size, input_dim].

        Returns:
            torch.Tensor: Projected embeddings in multimodal space.
                Shape: [batch_size, output_dim]
        """
        return self.mlp(x)

    def num_parameters(self) -> int:
        """
        Count trainable parameters in the projection head.

        Returns:
            int: Total number of trainable parameters. Should be << CLIP backbone size.
                For input_dim=768, hidden_dims=(512,256), output_dim=768:
                ~0.5M params vs 350M+ in full CLIP ViT-L/14
        """
        return sum(p.numel() for p in self.parameters())

    def num_floating_point_ops(self, input_shape: Tuple[int]) -> int:
        """
        Estimate floating point operations per forward pass.

        Args:
            input_shape (tuple): Shape of input tensor as (batch_size, input_dim).

        Returns:
            int: Estimated FLOPs for the entire MLP.
        """
        batch_size, in_dim = input_shape
        total_flops = 0

        current_dim = self.input_dim
        for hidden_dim in self.hidden_dims:
            # Linear layer: 2 * in_features * out_features (weights + bias)
            flops = 2 * current_dim * hidden_dim * batch_size
            total_flops += flops
            current_dim = hidden_dim

        # Final linear projection
        flops = 2 * current_dim * self.output_dim * batch_size
        total_flops += flops

        return int(total_flops)


def create_projection_head(
    input_dim: int = 768,
    output_dim: Optional[int] = None,
    hidden_dims: Tuple[int, ...] = (512, 256),
    dropout_rate: float = 0.3
) -> ProjectionHead:
    """
    Factory function for creating projection heads with sensible defaults.

    Args:
        input_dim: Dimensionality of change features (CLIP embedding size).
            Default 768 matches CLIP ViT-L/14.
        output_dim: Output dimension. Defaults to input_dim.
        hidden_dims: MLP bottleneck layers. Larger = more capacity, slower training.
        dropout_rate: Regularization strength. 0.3 is a good starting point.

    Returns:
        ProjectionHead: Configured adapter network.

    Example usage for different scenarios:
        # Default (matches CLIP ViT-L/14)
        >>> head = create_projection_head()
        >>> print(head.num_parameters())  # ~500k params

        # More capacity for hard negative mining
        >>> head = create_projection_head(
        ...     hidden_dims=(1024, 512, 256),
        ...     dropout_rate=0.2
        ... )

        # Minimal adapter (faster training, less expressive)
        >>> head = create_projection_head(
        ...     hidden_dims=(256,),
        ...     output_dim=768
        ... )
    """
    return ProjectionHead(input_dim, output_dim, hidden_dims, dropout_rate)


def save_adapter(path: str, adapter: "ProjectionHead", meta: Dict[str, Any]) -> None:
    """Persist a trained adapter together with the metadata needed to rebuild
    an identical ``ProjectionHead`` (dims, feature mode, encoder it targets)."""
    import os
    os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
    torch.save({"state_dict": adapter.state_dict(), "meta": meta}, path)


def load_adapter(path: str, map_location="cpu") -> Tuple["ProjectionHead", Dict[str, Any]]:
    """Reconstruct a ``ProjectionHead`` from a checkpoint written by
    :func:`save_adapter`. Returns ``(adapter_in_eval_mode, meta)``."""
    # weights_only=True: the adapter path is derived from user CLI dataset/encoder
    # strings, and save_adapter only ever stores tensors + a primitive meta dict
    # (dims/feature_mode/names) — so the safe loader suffices and we refuse to
    # unpickle arbitrary objects from a .pt that happens to sit at that path.
    ckpt = torch.load(path, map_location=map_location, weights_only=True)
    meta = ckpt["meta"]
    adapter = ProjectionHead(
        input_dim=meta["input_dim"],
        output_dim=meta["output_dim"],
        hidden_dims=tuple(meta["hidden_dims"]),
        dropout_rate=meta.get("dropout_rate", 0.3),
    )
    adapter.load_state_dict(ckpt["state_dict"])
    adapter.eval()
    return adapter, meta


def adapter_path(
    dataset: str,
    encoder: str,
    color_mode: str = "rgb",
    train_split: str = "train",
    mode: str = "difference",
    models_dir: str = "models",
) -> Path:
    """Canonical on-disk path for a trained ``ProjectionHead`` adapter.

    Single source of truth for the tagged filename, shared by the producer
    (``scripts.run_pipeline``) and the consumer (``scripts.export_results``) so the
    two can never drift: a non-default ``color_mode`` / ``train_split`` / ``mode``
    appends ``_<color>`` / ``_<split>`` / ``_<mode>`` respectively (``rgb`` / ``train``
    / ``difference`` add no suffix, keeping the committed
    ``<ds>__<enc>[_<color>]__adapter.pt`` names valid).
    """
    from src.embeddings import color_tag  # local import: keep model import-cycle-free
    split_tag = "" if train_split == "train" else f"_{train_split}"
    mode_tag = "" if mode == "difference" else f"_{mode}"
    return (Path(models_dir)
            / f"{dataset}__{encoder}{color_tag(color_mode)}{split_tag}{mode_tag}__adapter.pt")