| import torch |
| import torch.nn as nn |
| import math |
|
|
|
|
| class DualStreamTransformer(nn.Module): |
| def __init__( |
| self, |
| vocab_size: int, |
| d_model: int = 768, |
| n_head: int = 8, |
| d_hid: int = 768, |
| num_encoder_layers: int = 5, |
| num_decoder_layers: int = 8, |
| dino_dim: int = 768, |
| dropout: float = 0.1, |
| ): |
| super().__init__() |
| self.vocab_size = vocab_size |
| self.d_model = d_model |
| self.n_head = n_head |
| self.d_hid = d_hid |
| self.num_encoder_layers = num_encoder_layers |
| self.num_decoder_layers = num_decoder_layers |
| self.dino_dim = dino_dim |
| self.dropout = dropout |
|
|
| self.text_embedding = self.SimpleTextEmbedding(vocab_size, d_model) |
| self.image_embedding = self.DinoImageEmbedding(dino_dim, d_model) |
|
|
| self.image_encoder = self.Encoder( |
| d_model, n_head, d_hid, num_encoder_layers, dropout |
| ) |
|
|
| self.decoder = self.MultimodalDecoder( |
| d_model, n_head, d_hid, num_decoder_layers, dropout |
| ) |
|
|
| self.output_layer = nn.Linear(d_model, vocab_size) |
|
|
| def forward( |
| self, input_ids, dino_embedding=None, padding_mask=None, use_image: bool = False |
| ): |
| embedded = self.text_embedding(input_ids) |
|
|
| if ( |
| use_image |
| and dino_embedding is not None |
| and not torch.all(dino_embedding == 0) |
| ): |
| image_embedded = self.image_embedding(dino_embedding) |
| image_encoded = self.image_encoder(image_embedded) |
| else: |
| image_encoded = None |
|
|
| seq_len = embedded.size(1) |
|
|
| tgt_mask = self.decoder.generate_square_subsequent_mask(seq_len).to( |
| embedded.device |
| ) |
|
|
| decoder_output = self.decoder( |
| tgt=embedded, |
| image_memory=image_encoded, |
| tgt_mask=tgt_mask, |
| tgt_key_padding_mask=padding_mask, |
| ) |
|
|
| output = self.output_layer(decoder_output) |
|
|
| return output |
|
|
| class SimpleTextEmbedding(nn.Module): |
| def __init__(self, vocab_size, d_model, max_len=128, dropout=0.1): |
| super().__init__() |
| self.token_embedding = nn.Embedding(vocab_size, d_model) |
| self.position_embedding = nn.Embedding(max_len, d_model) |
| self.layer_norm = nn.LayerNorm(d_model) |
| self.dropout = nn.Dropout(p=dropout) |
| self.d_model = d_model |
|
|
| def forward(self, x): |
| batch_size, seq_len = x.size() |
|
|
| positions = ( |
| torch.arange(seq_len, device=x.device) |
| .unsqueeze(0) |
| .expand(batch_size, seq_len) |
| ) |
| scale = math.sqrt(self.d_model) |
|
|
| token_emb = self.token_embedding(x) * scale |
| pos_emb = self.position_embedding(positions) |
|
|
| embeddings = self.dropout(token_emb + pos_emb) |
|
|
| return self.layer_norm(embeddings) |
|
|
| class DinoImageEmbedding(nn.Module): |
| def __init__(self, dino_dim, d_model): |
| super().__init__() |
| self.projection_layer = nn.Linear(dino_dim, d_model) |
|
|
| def forward(self, x): |
| return self.projection_layer(x.unsqueeze(1)) |
|
|
| class Encoder(nn.Module): |
| def __init__( |
| self, |
| d_model: int, |
| n_head: int, |
| d_hid: int, |
| n_layers: int, |
| dropout: float = 0.1, |
| ): |
| super().__init__() |
| encoder_layer = nn.TransformerEncoderLayer( |
| d_model, n_head, d_hid, dropout, activation="gelu", batch_first=True |
| ) |
| self.encoder = nn.TransformerEncoder(encoder_layer, n_layers) |
|
|
| def forward(self, src, src_mask=None, src_key_padding_mask=None): |
| return self.encoder(src, src_mask, src_key_padding_mask) |
|
|
| class DynamicGating(nn.Module): |
| def __init__(self, d_model: int, dropout: float = 0.1): |
| super().__init__() |
| self.gate_fc = nn.Linear(d_model * 2, d_model) |
| self.dropout = nn.Dropout(dropout) |
| self.layer_norm = nn.LayerNorm(d_model) |
|
|
| def forward(self, text_features, image_features): |
| if image_features is None: |
| return text_features |
|
|
| combined = torch.cat([text_features, image_features], dim=-1) |
| gate = torch.sigmoid(self.gate_fc(combined)) |
| fused = gate * text_features + (1 - gate) * image_features |
| fused = self.layer_norm(self.dropout(fused)) |
| return fused |
|
|
| class MultimodalDecoderLayer(nn.Module): |
| def __init__(self, d_model: int, n_head: int, d_hid: int, dropout: float = 0.1): |
| super().__init__() |
| self.self_attn = nn.MultiheadAttention( |
| d_model, n_head, dropout=dropout, batch_first=True |
| ) |
| self.cross_attn_txt_image = nn.MultiheadAttention( |
| d_model, n_head, dropout=dropout, batch_first=True |
| ) |
|
|
| self.norm1 = nn.LayerNorm(d_model) |
| self.norm2 = nn.LayerNorm(d_model) |
| self.norm3 = nn.LayerNorm(d_model) |
|
|
| self.dropout = nn.Dropout(dropout) |
|
|
| self.gate = DualStreamTransformer.DynamicGating(d_model, dropout) |
|
|
| self.ff = nn.Sequential( |
| nn.Linear(d_model, d_hid), |
| nn.GELU(), |
| nn.Dropout(dropout), |
| nn.Linear(d_hid, d_model), |
| nn.Dropout(dropout), |
| ) |
|
|
| def forward(self, tgt, image_memory, tgt_mask=None, tgt_key_padding_mask=None): |
| tgt_norm = self.norm1(tgt) |
| self_attn_output, _ = self.self_attn( |
| tgt_norm, |
| tgt_norm, |
| tgt_norm, |
| key_padding_mask=tgt_key_padding_mask, |
| attn_mask=tgt_mask, |
| is_causal=True, |
| ) |
|
|
| tgt = tgt + self.dropout(self_attn_output) |
|
|
| if image_memory is not None: |
| tgt_norm = self.norm2(tgt) |
| cross_attn_output, _ = self.cross_attn_txt_image( |
| tgt_norm, image_memory, image_memory |
| ) |
| cross_attn_output = self.dropout(cross_attn_output) |
|
|
| fused = self.gate(tgt_norm, cross_attn_output) |
| tgt = tgt + fused |
|
|
| tgt_norm = self.norm3(tgt) |
| ff_output = self.ff(tgt_norm) |
| tgt = tgt + self.dropout(ff_output) |
|
|
| return tgt |
|
|
| class MultimodalDecoder(nn.Module): |
| def __init__( |
| self, |
| d_model: int, |
| n_head: int, |
| d_hid: int, |
| n_layers: int, |
| dropout: float = 0.1, |
| ): |
| super().__init__() |
| self.layers = nn.ModuleList( |
| [ |
| DualStreamTransformer.MultimodalDecoderLayer( |
| d_model, n_head, d_hid, dropout |
| ) |
| for _ in range(n_layers) |
| ] |
| ) |
|
|
| def generate_square_subsequent_mask(self, size): |
| mask = torch.triu(torch.ones(size, size), diagonal=1).bool() |
| return mask |
|
|
| def forward(self, tgt, image_memory, tgt_mask, tgt_key_padding_mask=None): |
| output = tgt |
| for layer in self.layers: |
| output = layer( |
| output, |
| image_memory, |
| tgt_mask=tgt_mask, |
| tgt_key_padding_mask=tgt_key_padding_mask, |
| ) |
| return output |
|
|
