Upload model.py with huggingface_hub

model.py

@@ -0,0 +1,229 @@
+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