model.py

"""MLP VAE for binarized MNIST with separate HF encoder/decoder submodules."""

from __future__ import annotations

from dataclasses import dataclass

import torch
import torch.nn as nn
from transformers import PreTrainedModel
from transformers.utils import ModelOutput

from .config import MNISTVAEConfig


@dataclass
class MNISTVAEOutput(ModelOutput):
    """Training forward outputs from :class:`MNISTVAE`."""

    logits: torch.Tensor | None = None
    target: torch.Tensor | None = None
    mean: torch.Tensor | None = None
    log_std: torch.Tensor | None = None


def _build_mlp(in_dim: int, hidden_dim: int, out_dim: int | None = None) -> nn.Sequential:
    layers: list[nn.Module] = [
        nn.Linear(in_dim, hidden_dim),
        nn.ReLU(),
        nn.Linear(hidden_dim, hidden_dim),
        nn.ReLU(),
    ]
    if out_dim is not None:
        layers.append(nn.Linear(hidden_dim, out_dim))
    return nn.Sequential(*layers)


class MNISTVAEEncoder(PreTrainedModel):
    """MLP encoder for binarized MNIST. Independently loadable via ``from_pretrained``."""

    config_class = MNISTVAEConfig

    def __init__(self, config: MNISTVAEConfig):
        super().__init__(config)
        self.encoder = _build_mlp(config.input_dim, config.hidden_dim)
        self.mean_linear = nn.Linear(config.hidden_dim, config.latent_dim)
        self.log_std_linear = nn.Linear(config.hidden_dim, config.latent_dim)
        self.post_init()

    def forward(self, pixel_values: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        hidden = self.encoder(pixel_values)
        return self.mean_linear(hidden), self.log_std_linear(hidden)

    def encode(self, pixel_values: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        return self.forward(pixel_values)


class MNISTVAEDecoder(PreTrainedModel):
    """MLP decoder for binarized MNIST. Independently loadable via ``from_pretrained``."""

    config_class = MNISTVAEConfig

    def __init__(self, config: MNISTVAEConfig):
        super().__init__(config)
        self.decoder = _build_mlp(config.latent_dim, config.hidden_dim, config.input_dim)
        self.post_init()

    def forward(self, z: torch.Tensor) -> torch.Tensor:
        return self.decoder(z)

    def forward_latent_positions(self, z: torch.Tensor) -> torch.Tensor:
        """Decode latent vectors to Bernoulli logits per pixel.

        Returns:
            Logits ``[batch, input_dim]`` (MuTAng-compatible decode entry point).
        """
        return self.forward(z)


def _sample_gaussian(mean: torch.Tensor, std: torch.Tensor) -> torch.Tensor:
    eps = torch.randn_like(mean)
    return mean + eps * std


class MNISTVAE(PreTrainedModel):
    """Composite MLP VAE for binarized MNIST training and Hub export."""

    config_class = MNISTVAEConfig
    base_model_prefix = "mnistvae"

    def __init__(self, config: MNISTVAEConfig):
        super().__init__(config)
        self.encoder = MNISTVAEEncoder(config)
        self.decoder = MNISTVAEDecoder(config)
        self.post_init()

    def forward(self, pixel_values: torch.Tensor, **kwargs) -> MNISTVAEOutput:
        del kwargs
        mean, log_std = self.encoder.encode(pixel_values)
        z = _sample_gaussian(mean, torch.exp(log_std))
        logits = self.decoder.forward_latent_positions(z)
        return MNISTVAEOutput(
            logits=logits,
            target=pixel_values,
            mean=mean,
            log_std=log_std,
        )

    def forward_latent_positions(self, z: torch.Tensor) -> torch.Tensor:
        """Decode latents to Bernoulli logits (HydrAMP-style API for MuTAng)."""
        return self.decoder.forward_latent_positions(z)