| """Inference-only unified-attention Forge2Vec architecture.""" |
|
|
| from __future__ import annotations |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from transformers import AutoConfig, AutoModel, SiglipVisionConfig, SiglipVisionModel |
|
|
|
|
| def _mean_pooling(tokens: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: |
| expanded = mask.unsqueeze(-1).to(tokens.dtype) |
| return (tokens * expanded).sum(1) / expanded.sum(1).clamp(min=1e-9) |
|
|
|
|
| class UnifiedAttentionForge2Vec(nn.Module): |
| """One 256-dimensional embedding model for text, metadata, and poster style.""" |
|
|
| def __init__(self, text_config_path: str, max_style_weight: float = 0.10) -> None: |
| super().__init__() |
| self.transformer = AutoModel.from_config(AutoConfig.from_pretrained(text_config_path)) |
| self.vision_encoder = SiglipVisionModel(SiglipVisionConfig()) |
| self.genre_embedding = nn.Embedding(100, 64, padding_idx=0) |
| self.style_bottleneck = nn.Sequential( |
| nn.Linear(4608, 256), nn.LayerNorm(256), nn.GELU(), |
| nn.Dropout(0.20), nn.Linear(256, 64), nn.LayerNorm(64), |
| ) |
| self.unified_description = nn.Sequential( |
| nn.Linear(768, 384), nn.LayerNorm(384), nn.GELU(), |
| nn.Dropout(0.10), nn.Linear(384, 128), |
| ) |
| self.unified_title = nn.Sequential(nn.Linear(384, 128), nn.LayerNorm(128), nn.GELU()) |
| self.unified_genre = nn.Sequential(nn.Linear(64, 128), nn.LayerNorm(128), nn.GELU()) |
| self.unified_metadata = nn.Sequential( |
| nn.Linear(4, 64), nn.LayerNorm(64), nn.GELU(), nn.Linear(64, 128), |
| ) |
| self.unified_style = nn.Sequential(nn.Linear(64, 128), nn.LayerNorm(128), nn.GELU()) |
| self.unified_cls = nn.Parameter(torch.empty(1, 1, 128)) |
| self.unified_missing = nn.Parameter(torch.empty(5, 128)) |
| self.unified_type = nn.Parameter(torch.empty(6, 128)) |
| semantic_layer = nn.TransformerEncoderLayer( |
| d_model=128, nhead=8, dim_feedforward=512, dropout=0.10, |
| activation="gelu", batch_first=True, norm_first=True, |
| ) |
| self.semantic_attention = nn.TransformerEncoder( |
| semantic_layer, num_layers=2, norm=nn.LayerNorm(128) |
| ) |
| fusion_layer = nn.TransformerEncoderLayer( |
| d_model=128, nhead=8, dim_feedforward=512, dropout=0.10, |
| activation="gelu", batch_first=True, norm_first=True, |
| ) |
| self.fusion_attention = nn.TransformerEncoder( |
| fusion_layer, num_layers=2, norm=nn.LayerNorm(128) |
| ) |
| self.image_gate = nn.Sequential( |
| nn.Linear(512, 128), nn.LayerNorm(128), nn.GELU(), |
| nn.Dropout(0.10), nn.Linear(128, 1), |
| ) |
| self.unified_output = nn.Sequential(nn.Linear(128, 256), nn.LayerNorm(256)) |
| self.register_buffer("max_style_weight", torch.tensor(float(max_style_weight))) |
|
|
| def encode_text(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor: |
| output = self.transformer(input_ids=input_ids, attention_mask=attention_mask) |
| |
| |
| |
| return _mean_pooling(output.last_hidden_state, attention_mask) |
|
|
| @staticmethod |
| def _present_or_missing(token: torch.Tensor, present: torch.Tensor, missing: torch.Tensor) -> torch.Tensor: |
| present = present.reshape(-1, 1, 1) |
| return token * present + missing.reshape(1, 1, -1) * (1.0 - present) |
|
|
| def _semantic_tokens(self, ua, en, titles, genres, metadata, mask): |
| batch_size = ua.size(0) |
| description = self.unified_description(torch.cat((ua, en), -1)).unsqueeze(1) |
| title = self.unified_title(titles).unsqueeze(1) |
| genre_vectors = self.genre_embedding(genres.long()) |
| genre_mask = (genres != 0).unsqueeze(-1).to(genre_vectors.dtype) |
| genre_pooled = (genre_vectors * genre_mask).sum(1) / genre_mask.sum(1).clamp(min=1.0) |
| genre = self.unified_genre(genre_pooled).unsqueeze(1) |
| metadata_token = self.unified_metadata(metadata).unsqueeze(1) |
| raw = (description, title, genre, metadata_token) |
| tokens = [ |
| self._present_or_missing(token, mask[:, index], self.unified_missing[index]) |
| for index, token in enumerate(raw) |
| ] |
| sequence = torch.cat((self.unified_cls.expand(batch_size, -1, -1), *tokens), dim=1) |
| sequence = sequence + self.unified_type[:5].unsqueeze(0) |
| encoded = self.semantic_attention(sequence) |
| return encoded, encoded[:, 0] |
|
|
| def encode_style(self, pixels: torch.Tensor) -> torch.Tensor: |
| output = self.vision_encoder(pixel_values=pixels, output_hidden_states=True) |
| statistics = [] |
| for tokens in (output.hidden_states[4], output.hidden_states[8], output.hidden_states[-1]): |
| statistics.extend((tokens.mean(1), tokens.std(1, unbiased=False))) |
| return F.normalize(self.style_bottleneck(torch.cat(statistics, -1)), dim=-1) |
|
|
| def forward(self, ua, en, titles, genres, metadata, modality_mask, poster_pixels=None): |
| semantic_sequence, semantic_summary = self._semantic_tokens( |
| ua, en, titles, genres, metadata, modality_mask |
| ) |
| available = modality_mask[:, 4:5] |
| image_token = self.unified_missing[4].reshape(1, 128).expand(ua.size(0), -1) |
| if poster_pixels is not None and bool((available > 0).any()): |
| active = available.squeeze(1) > 0 |
| image_token = image_token.clone() |
| image_token[active] = self.unified_style(self.encode_style(poster_pixels[active])) |
| gate_features = torch.cat(( |
| semantic_summary, image_token, torch.abs(semantic_summary - image_token), |
| semantic_summary * image_token, |
| ), -1) |
| gate = torch.sigmoid(self.image_gate(gate_features)) * available * self.max_style_weight |
| sequence = torch.cat((semantic_sequence, (image_token * gate).unsqueeze(1)), 1) |
| fused = self.fusion_attention(sequence + self.unified_type.unsqueeze(0)) |
| return F.normalize(self.unified_output(fused[:, 0]), dim=-1) |
|
|
