File size: 7,930 Bytes
b233cf7 4888d21 b233cf7 cb7920d b233cf7 cb7920d b233cf7 4888d21 b233cf7 4888d21 8ce6ce2 b233cf7 4888d21 b233cf7 4888d21 cb7920d 254d9db 4888d21 254d9db 4888d21 b233cf7 4888d21 b233cf7 8ce6ce2 4888d21 b233cf7 4888d21 8ce6ce2 b233cf7 4888d21 cb7920d b233cf7 8ce6ce2 254d9db 8ce6ce2 b233cf7 4888d21 254d9db b233cf7 254d9db 4888d21 254d9db b233cf7 8ce6ce2 254d9db 8ce6ce2 b233cf7 254d9db b233cf7 8ce6ce2 b233cf7 8ce6ce2 b233cf7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 | 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
|