Upload model.py with huggingface_hub

model.py

@@ -3,14 +3,13 @@ import torch.nn as nn
 import torch.nn.functional as F
 from typing import Dict
 
-# ... Paste the full code for TextFieldAttention, GenreSelfAttention, and AnimeEmbeddingModel here ...
-# (The same code as in Cell 2 of the notebook)
-
 class TextFieldAttention(nn.Module):
+    """Calculates a weighted sum of text field embeddings."""
     def __init__(self, num_fields: int, field_dim: int):
         super().__init__()
         self.attn = nn.Linear(field_dim, 1, bias=False)
         self.num_fields = num_fields
+
     def forward(self, fields: torch.Tensor):
         scores = self.attn(fields)
         weights = F.softmax(scores, dim=1)
@@ -18,9 +17,11 @@ class TextFieldAttention(nn.Module):
         return weighted_sum, weights.squeeze(-1)
 
 class GenreSelfAttention(nn.Module):
+    """Calculates a weighted sum of genres based only on the genres themselves."""
     def __init__(self, genre_dim: int):
         super().__init__()
         self.attn_scorer = nn.Linear(genre_dim, 1, bias=False)
+
     def forward(self, genre_embeds: torch.Tensor, mask: torch.Tensor):
         scores = self.attn_scorer(genre_embeds)
         scores.masked_fill_(mask == 0, -1e9)
@@ -28,46 +29,94 @@ class GenreSelfAttention(nn.Module):
         weighted_sum = (genre_embeds * weights).sum(dim=1)
         return weighted_sum
 
+class ModalityAttention(nn.Module):
+    """
+    Calculates a weighted sum of vectors from different modalities (text, genres, etc.),
+    allowing the model to dynamically determine their importance.
+    """
+    def __init__(self, num_modalities: int, modality_dim: int):
+        super().__init__()
+        self.attn_scorer = nn.Linear(modality_dim, 1, bias=False)
+        self.num_modalities = num_modalities
+
+    def forward(self, modalities: torch.Tensor):
+        scores = self.attn_scorer(modalities)
+        weights = F.softmax(scores, dim=1)
+        weighted_sum = (modalities * weights).sum(dim=1)
+        return weighted_sum, weights.squeeze(-1)
+
 class AnimeEmbeddingModel(nn.Module):
-    def __init__(self, vocab_sizes: Dict[str, int], embedding_dims: Dict[str, int], dropout_rate: float = 0.3, text_embedding_size: int = 384):
+    """
+    Main model v13.
+    """
+    def __init__(self,
+                 vocab_sizes: Dict[str, int],
+                 embedding_dims: Dict[str, int] = None,
+                 dropout_rate: float = 0.3,
+                 text_embedding_size: int = 384,
+                 final_embedding_dim: int = 512):
         super().__init__()
+
+        if embedding_dims is None:
+            embedding_dims = {'genre': 128, 'studio': 64, 'type': 16, 'numerical': 32}
+
         self.embedding_dims = embedding_dims
+        self.final_embedding_dim = final_embedding_dim
+
         self.genre_embedding = nn.Embedding(vocab_sizes['genre'], embedding_dims['genre'], padding_idx=0)
         self.studio_embedding = nn.Embedding(vocab_sizes['studio'], embedding_dims['studio'])
         self.type_embedding = nn.Embedding(vocab_sizes['type'], embedding_dims['type'])
         self.numerical_layer = nn.Linear(6, embedding_dims['numerical'])
+
+        self.text_projector = nn.Linear(text_embedding_size, final_embedding_dim)
+        self.genre_projector = nn.Linear(embedding_dims['genre'], final_embedding_dim)
+        other_dim = embedding_dims['studio'] + embedding_dims['type'] + embedding_dims['numerical']
+        self.other_projector = nn.Linear(other_dim, final_embedding_dim)
+
         self.text_field_attention = TextFieldAttention(num_fields=6, field_dim=text_embedding_size)
         self.genre_attention = GenreSelfAttention(embedding_dims['genre'])
-        total_dim = sum(embedding_dims.values())
+        self.modality_attention = ModalityAttention(num_modalities=3, modality_dim=final_embedding_dim)
+
         self.encoder = nn.Sequential(
-            nn.Linear(total_dim, 1024), nn.ReLU(), nn.Dropout(dropout_rate), nn.LayerNorm(1024),
+            nn.Linear(final_embedding_dim, 1024), nn.ReLU(), nn.Dropout(dropout_rate), nn.LayerNorm(1024),
             nn.Linear(1024, 768), nn.ReLU(), nn.Dropout(dropout_rate), nn.LayerNorm(768),
-            nn.Linear(768, 512),
+            nn.Linear(768, final_embedding_dim),
         )
+
         self.text_scale = nn.Parameter(torch.tensor(1.0))
         self.genre_scale = nn.Parameter(torch.tensor(1.0))
         self.other_scale = nn.Parameter(torch.tensor(1.0))
+
     def forward(self, batch: Dict[str, torch.Tensor]) -> torch.Tensor:
         text_fields = torch.stack([
             batch['precomputed_ua_desc'], batch['precomputed_en_desc'],
             batch['precomputed_ua_title'], batch['precomputed_en_title'],
             batch['precomputed_original_title'], batch['precomputed_alternate_names'],
         ], dim=1)
-        text_vector, _ = self.text_field_attention(text_fields)
+        text_vector_raw, _ = self.text_field_attention(text_fields)
+
         genre_embeds_raw = self.genre_embedding(batch['genres'])
-        genre_vector = self.genre_attention(genre_embeds_raw, batch['genres_mask'].unsqueeze(-1))
+        genre_vector_raw = self.genre_attention(genre_embeds_raw, batch['genres_mask'].unsqueeze(-1))
+
         studio_emb = self.studio_embedding(batch['studio'])
         type_emb = self.type_embedding(batch['type'])
         numerical_emb = F.relu(self.numerical_layer(batch['numerical']))
         other_vector_parts = torch.cat([studio_emb, type_emb, numerical_emb], dim=1)
-        text_vector_norm = F.normalize(text_vector, p=2, dim=1)
-        genre_vector_norm = F.normalize(genre_vector, p=2, dim=1)
-        other_vector_norm = F.normalize(other_vector_parts, p=2, dim=1)
-        scaled_text = text_vector_norm * self.text_scale
-        scaled_genre = genre_vector_norm * self.genre_scale
-        scaled_other = other_vector_norm * self.other_scale
-        combined = torch.cat([scaled_text, scaled_genre, scaled_other], dim=1)
+
+        text_vector_proj = self.text_projector(text_vector_raw)
+        genre_vector_proj = self.genre_projector(genre_vector_raw)
+        other_vector_proj = self.other_projector(other_vector_parts)
+
+        modalities = torch.stack([
+            F.normalize(text_vector_proj, p=2, dim=1) * self.text_scale,
+            F.normalize(genre_vector_proj, p=2, dim=1) * self.genre_scale,
+            F.normalize(other_vector_proj, p=2, dim=1) * self.other_scale,
+        ], dim=1)
+
+        combined, _ = self.modality_attention(modalities)
+
         embedding_logits = self.encoder(combined)
         embedding = torch.tanh(embedding_logits)
         embedding = F.normalize(embedding, p=2, dim=1)
+
         return embedding