feat: Upload anime2vec v12 model and artifacts

README.md

@@ -4,12 +4,45 @@ language:
 - en
 - uk
 - ja
-library_name: pytorch
 tags:
 - anime
-- recommender-system
-- sentence-similarity
-- metric-learning
-- feature-extraction
-pipeline_tag: feature-extraction
----
+- embeddings
+- semantic-search
+- vector-arithmetic
+- pytorch
+datasets:
+- private
+author: Lorg0n
+---
+
+# hikka-forge-anime2vec
+
+This repository contains `hikka-forge-anime2vec`, a sophisticated semantic vector space model for anime, created by [Lorg0n](https://huggingface.co/Lorg0n).
+
+The model is trained to understand deep connections between titles based on multilingual textual descriptions, genres, studios, and other metadata. It supports vector arithmetic, allowing for creative queries like `"Show me something like 'Spirited Away' - 'Ghibli Style' + 'Cyberpunk'"`.
+
+## Model Details
+
+- **Model Version**: v12
+- **Architecture**: A multi-input neural network with separate processing streams for text, genres, and other categorical/numerical features. It uses attention mechanisms to weigh the importance of different text fields and genres.
+- **Training**: Trained using a combination of Triplet Loss (from explicit user recommendations), Cosine Similarity Loss for vector arithmetic examples, and a Diversity Loss to ensure a well-distributed embedding space.
+- **Data**: Trained on a private, non-public database of anime titles.
+
+## How to Use
+
+*This model requires custom code for loading and inference due to its unique architecture and preprocessing steps.*
+
+A full usage example will be provided soon. The general workflow involves:
+1. Loading the model, config, and pickled `LabelEncoder` objects.
+2. Preprocessing new anime data (fetching from a data source, encoding text with a SentenceTransformer, etc.).
+3. Using the model to generate a 512-dimensional embedding.
+4. Performing similarity search or vector arithmetic in the embedding space.
+
+## Files in this Repository
+
+This repository contains all files necessary for model inference:
+
+- `pytorch_model.bin`: The trained model weights.
+- `config.json`: Configuration file specifying model architecture and vocabulary sizes.
+- `model.py`: The Python code defining the `AnimeEmbeddingModel` class.
+- `le_genre.pkl`, `le_studio.pkl`, `le_type.pkl`: Pickled Scikit-learn `LabelEncoder` objects required for preprocessing new data.

config.json

@@ -0,0 +1,19 @@
+{
+    "architecture": "AnimeEmbeddingModel",
+    "vocab_sizes": {
+        "genre": 77,
+        "studio": 1170,
+        "type": 7
+    },
+    "embedding_dims": {
+        "genre": 128,
+        "studio": 64,
+        "type": 16,
+        "numerical": 32,
+        "text": 384
+    },
+    "text_embedding_size": 384,
+    "model_type": "hikka-forge-anime2vec",
+    "model_version": "v12",
+    "author": "Lorg0n"
+}

le_genre.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5b281e8d281d95b8a3681c145b23c74dd916949da2b881697c3053927d00bb8b
+size 5476

le_studio.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:82be27d7b828b66ed776e59f9e57093e461359eb7b6d783cecbb45697c905785
+size 229570

le_type.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:11535e48ef762e4f790f0013a47477c258a3a6dd0cfaa45dfc0e48557a9411ab
+size 432

model.py

@@ -0,0 +1,70 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Dict
+
+class TextFieldAttention(nn.Module):
+    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)
+        weighted_sum = (fields * weights).sum(dim=1)
+        return weighted_sum, weights.squeeze(-1)
+
+class GenreSelfAttention(nn.Module):
+    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)
+        weights = F.softmax(scores, dim=1)
+        weighted_sum = (genre_embeds * weights).sum(dim=1)
+        return weighted_sum
+
+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):
+        super().__init__()
+        self.embedding_dims = embedding_dims
+        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_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.encoder = nn.Sequential(
+            nn.Linear(total_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),
+        )
+        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)
+        genre_embeds_raw = self.genre_embedding(batch['genres'])
+        genre_vector = 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)
+        embedding_logits = self.encoder(combined)
+        embedding = torch.tanh(embedding_logits)
+        embedding = F.normalize(embedding, p=2, dim=1)
+        return embedding

pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:208dab2e050ff92356d680320ad10631afa5f5ff9c7ba605f519a944d15c3690
+size 7647786