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