src/pino/pimt_model_hf.py

from __future__ import annotations

from typing import Optional, Tuple, Union

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import PretrainedConfig, PreTrainedModel
from transformers.modeling_outputs import SequenceClassifierOutput
from pino.pimt_model import DEFAULT_EMBEDDING_DIM
from pino.heads import PIMTHeads


class PIMTConfig(PretrainedConfig):
    """
    Hugging Face PretrainedConfig for the Physics-Informed Mixture Transformer.
    """

    model_type = "pimt"

    def __init__(
        self,
        embedding_dim: int = DEFAULT_EMBEDDING_DIM,
        state_dim: int = 2,
        hidden_dim: int = 256,
        num_heads: int = 8,
        num_layers: int = 4,
        dropout: float = 0.1,
        max_ingredients: int = 32,
        num_classes_sub: int = 7,
        text_embedding_dim: int = 384,
        **kwargs,
    ) -> None:
        self.embedding_dim = embedding_dim
        self.state_dim = state_dim
        self.hidden_dim = hidden_dim
        self.num_heads = num_heads
        self.num_layers = num_layers
        self.dropout = dropout
        self.max_ingredients = max_ingredients
        self.num_classes_sub = num_classes_sub
        self.text_embedding_dim = text_embedding_dim
        kwargs.setdefault("return_dict", True)
        super().__init__(**kwargs)


class PhysicsInformedGatingBlock(nn.Module):
    """
    Modulate a static token embedding with a time-varying physical state vector.
    """

    def __init__(self, embedding_dim: int, state_dim: int, hidden_dim: int = 64) -> None:
        super().__init__()
        self.embedding_dim = embedding_dim
        self.state_dim = state_dim
        self.gate = nn.Sequential(
            nn.Linear(embedding_dim + state_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, embedding_dim),
            nn.Sigmoid(),
        )

    def forward(self, tokens: torch.Tensor, states: torch.Tensor) -> torch.Tensor:
        b, seq_len, e = tokens.shape
        t_steps = states.size(1)
        tokens_t = tokens.unsqueeze(1).expand(b, t_steps, seq_len, e)
        x = torch.cat([tokens_t, states], dim=-1)
        gate = self.gate(x)
        return tokens_t * gate


class PhysicsInformedMixtureTransformer(PreTrainedModel):
    """
    Text-controllable, permutation-invariant transformer encoder for fragrance
    dry-down trajectories. Accepts a 384-D sentence-transformer embedding to
    gate the trajectory representation.
    """

    config_class = PIMTConfig
    base_model_prefix = "pimt"
    supports_gradient_checkpointing = False

    def __init__(self, config: PIMTConfig) -> None:
        super().__init__(config)
        self.embedding_dim = config.embedding_dim
        self.hidden_dim = config.hidden_dim

        self.gating = PhysicsInformedGatingBlock(
            config.embedding_dim, config.state_dim, hidden_dim=config.hidden_dim
        )
        self.input_proj = nn.Linear(config.embedding_dim, config.hidden_dim)

        encoder_layer = nn.TransformerEncoderLayer(
            d_model=config.hidden_dim,
            nhead=config.num_heads,
            dim_feedforward=config.hidden_dim * 4,
            dropout=config.dropout,
            batch_first=True,
        )
        self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=config.num_layers)
        self.text_proj = nn.Linear(config.text_embedding_dim, config.hidden_dim)
        # Objective output vocabulary stays 138-D; structural input is 151-D.
        objective_dim = getattr(config, "objective_dim", 138)
        self.heads = PIMTHeads(
            hidden_dim=config.hidden_dim,
            objective_dim=objective_dim,
            num_classes_sub=getattr(config, "num_classes_sub", 7),
        )

        self.post_init()

    def _compute_multitask_loss(
        self,
        pred_obj: torch.Tensor,
        target_obj: torch.Tensor,
        pred_sub: torch.Tensor,
        target_sub: torch.Tensor,
    ) -> torch.Tensor:
        """Compute a weighted multi-task loss: MSE for objective + MSE for 7-D psychometric vector."""
        weights = {"obj": 1.0, "sub": 0.5}

        obj_mask = (target_obj.sum(dim=-1) > 0).float().unsqueeze(-1)
        obj_loss = F.mse_loss(pred_obj, target_obj, reduction="none")
        obj_loss = (obj_loss * obj_mask).sum() / obj_mask.sum().clamp_min(1.0)

        if target_sub.dim() == 3:
            target_sub = target_sub.mean(dim=1)
        sub_loss = F.mse_loss(pred_sub, target_sub)

        total = weights["obj"] * obj_loss + weights["sub"] * sub_loss
        return total

    def forward(
        self,
        tokens: torch.Tensor,
        physics: torch.Tensor,
        src_key_padding_mask: Optional[torch.Tensor] = None,
        text_embedding: Optional[torch.Tensor] = None,
        labels: Optional[torch.Tensor] = None,
        labels_obj: Optional[torch.Tensor] = None,
        labels_sub: Optional[torch.Tensor] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[tuple, SequenceClassifierOutput]:
        """
        tokens: (B, S, E)
        physics: (B, T, S, state_dim)
        src_key_padding_mask: (B, S) bool, True for padding positions.
        text_embedding: (B, 384) sentence-transformer embedding for conditioning.
        labels: fused (B, T, 151) tensor of objective + subjective labels (legacy HF shape).
        """
        return_dict = return_dict if return_dict is not None else self.config.return_dict

        if labels is not None:
            labels_obj = labels[..., :138]
            labels_sub = labels[..., 138:]

        gated = self.gating(tokens, physics)  # (B, T, S, E)
        b, t, s, e = gated.shape
        x = self.input_proj(gated)  # (B, T, S, H)

        x = x.reshape(b * t, s, self.hidden_dim)
        if src_key_padding_mask is not None:
            mask_bt = src_key_padding_mask.unsqueeze(1).expand(-1, t, -1).reshape(b * t, s)
            x = self.encoder(x, src_key_padding_mask=mask_bt)
        else:
            x = self.encoder(x)
        x = x.reshape(b, t, s, self.hidden_dim)

        # Project text embedding and use it as a global conditioning gate.
        if text_embedding is not None:
            text_hidden = self.text_proj(text_embedding)  # (B, H)
            text_gate = F.sigmoid(text_hidden).view(b, 1, 1, self.hidden_dim)  # (B, 1, 1, H)
            x = x * text_gate

        logits = self.heads(x)  # {"objective": (B, T, 138), "subjective": (B, 7), "alignment": (B, 138)}

        # Fuse objective trajectory + subjective + alignment into a single tensor
        # for HF Trainer labels (obj T x 138, then subjective 7, then alignment 138).
        subjective_expanded = logits["subjective"].unsqueeze(1).expand(-1, t, -1)
        alignment_expanded = logits["alignment"].unsqueeze(1).expand(-1, t, -1)
        fused_logits = torch.cat([logits["objective"], subjective_expanded, alignment_expanded], dim=-1)

        # Labels are fused (B, T, 138 + 7 + 138) = (B, T, 283) in HF Trainer.
        if labels is not None:
            labels_obj = labels[..., :138]
            labels_sub = labels[..., 138:145]
            labels_alignment = labels[..., 145:283]

        loss = None
        if labels_obj is not None and labels_sub is not None and labels_alignment is not None:
            alignment_target = labels_alignment.mean(dim=1)  # (B, 138)
            alignment_loss = 1.0 - F.cosine_similarity(logits["alignment"], alignment_target, dim=-1).mean()
            loss = self._compute_multitask_loss(
                logits["objective"], labels_obj, logits["subjective"], labels_sub
            ) + alignment_loss

        output = SequenceClassifierOutput(
            loss=loss,
            logits=fused_logits,
            hidden_states=None,
            attentions=None,
        )

        if not return_dict:
            return (output.loss, output.logits) if output.loss is not None else (output.logits,)

        return output