File size: 6,225 Bytes
13065f0 | 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 | """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)
# Training and production inference use SentenceTransformer's raw
# mean-pooled vectors here. Normalization happens only at the final
# Forge2Vec output.
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)
|