Upload 41 files

models/__init__.py

@@ -0,0 +1,14 @@
+from .bert import BERTModel
+from .dae import DAEModel
+from .vae import VAEModel
+
+MODELS = {
+    BERTModel.code(): BERTModel,
+    DAEModel.code(): DAEModel,
+    VAEModel.code(): VAEModel
+}
+
+
+def model_factory(args):
+    model = MODELS[args.model_code]
+    return model(args)

models/base.py

@@ -0,0 +1,15 @@
+import torch.nn as nn
+
+from abc import *
+
+
+class BaseModel(nn.Module, metaclass=ABCMeta):
+    def __init__(self, args):
+        super().__init__()
+        self.args = args
+
+    @classmethod
+    @abstractmethod
+    def code(cls):
+        pass
+

models/bert.py

@@ -0,0 +1,19 @@
+from .base import BaseModel
+from .bert_modules.bert import BERT
+
+import torch.nn as nn
+
+
+class BERTModel(BaseModel):
+    def __init__(self, args):
+        super().__init__(args)
+        self.bert = BERT(args)
+        self.out = nn.Linear(self.bert.hidden, args.num_items + 1)
+
+    @classmethod
+    def code(cls):
+        return 'bert'
+
+    def forward(self, x):
+        x = self.bert(x)
+        return self.out(x)

models/bert_modules/__init__.py

@@ -0,0 +1 @@
+

models/bert_modules/attention/__init__.py

@@ -0,0 +1,2 @@
+from .multi_head import MultiHeadedAttention
+from .single import Attention

models/bert_modules/attention/multi_head.py

@@ -0,0 +1,37 @@
+import torch.nn as nn
+from .single import Attention
+
+
+class MultiHeadedAttention(nn.Module):
+    """
+    Take in model size and number of heads.
+    """
+
+    def __init__(self, h, d_model, dropout=0.1):
+        super().__init__()
+        assert d_model % h == 0
+
+        # We assume d_v always equals d_k
+        self.d_k = d_model // h
+        self.h = h
+
+        self.linear_layers = nn.ModuleList([nn.Linear(d_model, d_model) for _ in range(3)])
+        self.output_linear = nn.Linear(d_model, d_model)
+        self.attention = Attention()
+
+        self.dropout = nn.Dropout(p=dropout)
+
+    def forward(self, query, key, value, mask=None):
+        batch_size = query.size(0)
+
+        # 1) Do all the linear projections in batch from d_model => h x d_k
+        query, key, value = [l(x).view(batch_size, -1, self.h, self.d_k).transpose(1, 2)
+                             for l, x in zip(self.linear_layers, (query, key, value))]
+
+        # 2) Apply attention on all the projected vectors in batch.
+        x, attn = self.attention(query, key, value, mask=mask, dropout=self.dropout)
+
+        # 3) "Concat" using a view and apply a final linear.
+        x = x.transpose(1, 2).contiguous().view(batch_size, -1, self.h * self.d_k)
+
+        return self.output_linear(x)

models/bert_modules/attention/single.py

@@ -0,0 +1,25 @@
+import torch.nn as nn
+import torch.nn.functional as F
+import torch
+
+import math
+
+
+class Attention(nn.Module):
+    """
+    Compute 'Scaled Dot Product Attention
+    """
+
+    def forward(self, query, key, value, mask=None, dropout=None):
+        scores = torch.matmul(query, key.transpose(-2, -1)) \
+                 / math.sqrt(query.size(-1))
+
+        if mask is not None:
+            scores = scores.masked_fill(mask == 0, -1e9)
+
+        p_attn = F.softmax(scores, dim=-1)
+
+        if dropout is not None:
+            p_attn = dropout(p_attn)
+
+        return torch.matmul(p_attn, value), p_attn

models/bert_modules/bert.py

@@ -0,0 +1,45 @@
+from torch import nn as nn
+
+from models.bert_modules.embedding import BERTEmbedding
+from models.bert_modules.transformer import TransformerBlock
+from utils import fix_random_seed_as
+
+from pathlib import Path
+import pickle
+class BERT(nn.Module):
+    def __init__(self, args):
+        super().__init__()
+
+        fix_random_seed_as(args.model_init_seed)
+        # self.init_weights()
+
+        max_len = args.bert_max_len
+        num_items = args.num_items
+        n_layers = args.bert_num_blocks
+        heads = args.bert_num_heads
+        vocab_size = num_items + 2
+        hidden = args.bert_hidden_units
+        self.hidden = hidden
+        dropout = args.bert_dropout
+
+        # embedding for BERT, sum of positional, segment, token embeddings
+        self.embedding = BERTEmbedding(vocab_size=vocab_size, embed_size=self.hidden, max_len=max_len, dropout=dropout)
+
+        # multi-layers transformer blocks, deep network
+        self.transformer_blocks = nn.ModuleList(
+            [TransformerBlock(hidden, heads, hidden * 4, dropout) for _ in range(n_layers)])
+
+    def forward(self, x):
+        mask = (x > 0).unsqueeze(1).repeat(1, x.size(1), 1).unsqueeze(1)
+
+        # embedding the indexed sequence to sequence of vectors
+        x = self.embedding(x)
+
+        # running over multiple transformer blocks
+        for transformer in self.transformer_blocks:
+            x = transformer.forward(x, mask)
+
+        return x
+
+    def init_weights(self):
+        pass

models/bert_modules/embedding/__init__.py

@@ -0,0 +1 @@
+from .bert import BERTEmbedding

models/bert_modules/embedding/bert.py

@@ -0,0 +1,304 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import pickle
+import json
+import threading
+from pathlib import Path
+import torch.nn as nn
+from .token import TokenEmbedding
+from .position import PositionalEmbedding
+import time
+from pathlib import Path
+from torch.nn.utils.rnn import pad_sequence
+import pickle 
+import json
+import os
+
+class BERTEmbedding(nn.Module):
+    _mappings_cache = None
+    _cache_lock = threading.Lock()
+    
+    @classmethod
+    def _load_mappings(cls):
+        if cls._mappings_cache is None:
+            with cls._cache_lock:
+                if cls._mappings_cache is None:  # Double-checked locking
+                    try:
+
+                        main_dir = os.getcwd()
+
+                        relative_path_dataset = "Data/preprocessed/AnimeRatings_min_rating7-min_uc10-min_sc10-splitleave_one_out/dataset.pkl"
+                        relative_path_genres = "Data/AnimeRatings/id_to_genreids.json"
+
+                        full_path_dataset = Path(main_dir) / relative_path_dataset
+                        full_path_genres = Path(main_dir) / relative_path_genres
+
+
+                        with full_path_dataset.open('rb') as f:
+                            dataset_smap = pickle.load(f)["smap"]
+                        
+                        with full_path_genres.open('rb') as f:
+                            id_to_genres = json.load(f)
+                        
+                        cls._mappings_cache = {
+                            'dataset_smap': dataset_smap,
+                            'id_to_genres': id_to_genres
+                        }
+                        
+                    except Exception as e:
+                        print(f"Warning: Could not load mappings: {e}")
+                        cls._mappings_cache = {
+                            'dataset_smap': {},
+                            'id_to_genres': {}
+                        }
+        return cls._mappings_cache
+    
+    def __init__(self, vocab_size, embed_size, max_len, dropout=0.1, multi_genre=True, max_genres_per_anime=5):
+        super().__init__()
+
+        mappings = self._load_mappings()
+        dataset_smap = mappings['dataset_smap']
+        id_to_genres = mappings['id_to_genres']
+        
+        self.multi_genre = multi_genre
+        self.max_genres_per_anime = max_genres_per_anime
+
+        all_genres = set()
+        for anime_id, genres in id_to_genres.items():
+            all_genres.update(genres)
+
+        max_genre_id = max(all_genres) if all_genres else 0
+        self.num_genres = max_genre_id + 1
+        
+        print(f"Detected {self.num_genres} unique genres (max_id: {max_genre_id})")
+
+        self.vocab_size = vocab_size
+        
+        if multi_genre:
+            self._create_multi_genre_mapping(dataset_smap, id_to_genres, vocab_size)
+        else:
+            self._create_single_genre_mapping(dataset_smap, id_to_genres, vocab_size)
+
+        self.token = TokenEmbedding(vocab_size=vocab_size, embed_size=embed_size)
+        self.genre_embed = nn.Embedding(num_embeddings=self.num_genres, embedding_dim=embed_size, padding_idx=0)
+
+        if multi_genre:
+            self.fusion_layer = nn.Sequential(
+                nn.Linear(embed_size * 2, embed_size),
+                nn.LayerNorm(embed_size),
+                nn.ReLU()
+            )
+
+            self.genre_aggregation = nn.Parameter(torch.ones(max_genres_per_anime) / max_genres_per_anime)
+            self.genre_attention = nn.MultiheadAttention(embed_size, num_heads=4, batch_first=True)
+        else:
+            self.fusion_layer = nn.Sequential(
+                nn.Linear(embed_size * 2, embed_size),
+                nn.LayerNorm(embed_size),
+                nn.ReLU()
+            )
+        
+        self.dropout = nn.Dropout(p=dropout)
+        self.embed_size = embed_size
+
+        self._genre_cache = {}
+        self._cache_lock = threading.Lock()
+    
+    def _create_single_genre_mapping(self, dataset_smap, id_to_genres, vocab_size):
+        token_to_genre = {}
+        for anime_id, token_id in dataset_smap.items():
+            if token_id < vocab_size:
+                genre_list = id_to_genres.get(str(anime_id), [0])
+                genre_id = genre_list[0] if genre_list else 0
+                if genre_id >= self.num_genres:
+                    print(f"Warning: Genre ID {genre_id} >= {self.num_genres}, setting to 0")
+                    genre_id = 0
+                token_to_genre[token_id] = genre_id
+        
+        if token_to_genre:
+            token_ids = torch.tensor(list(token_to_genre.keys()), dtype=torch.long)
+            genre_ids = torch.tensor(list(token_to_genre.values()), dtype=torch.long)
+            
+            self.register_buffer('token_ids', token_ids)
+            self.register_buffer('genre_ids', genre_ids)
+            self.has_mappings = True
+        else:
+            self.register_buffer('token_ids', torch.empty(0, dtype=torch.long))
+            self.register_buffer('genre_ids', torch.empty(0, dtype=torch.long))
+            self.has_mappings = False
+    
+    def _create_multi_genre_mapping(self, dataset_smap, id_to_genres, vocab_size):
+        token_to_genres = {}
+        for anime_id, token_id in dataset_smap.items():
+            if token_id < vocab_size:
+                genre_list = id_to_genres.get(str(anime_id), [0])
+
+                valid_genres = []
+                for genre_id in genre_list:
+                    if genre_id >= self.num_genres:
+                        print(f"Warning: Genre ID {genre_id} >= {self.num_genres}, setting to 0")
+                        genre_id = 0
+                    valid_genres.append(genre_id)
+
+                if len(valid_genres) < self.max_genres_per_anime:
+                    valid_genres.extend([0] * (self.max_genres_per_anime - len(valid_genres)))
+                else:
+                    valid_genres = valid_genres[:self.max_genres_per_anime]
+                
+                token_to_genres[token_id] = valid_genres
+        
+        if token_to_genres:
+            token_ids = torch.tensor(list(token_to_genres.keys()), dtype=torch.long)
+            genre_ids = torch.tensor(list(token_to_genres.values()), dtype=torch.long)
+            
+            self.register_buffer('token_ids', token_ids)
+            self.register_buffer('genre_ids', genre_ids)
+            self.has_mappings = True
+        else:
+            self.register_buffer('token_ids', torch.empty(0, dtype=torch.long))
+            self.register_buffer('genre_ids', torch.empty(0, self.max_genres_per_anime, dtype=torch.long))
+            self.has_mappings = False
+    
+    def _get_single_genre_mapping(self, sequence):
+        """Original single genre mapping with improved bounds checking"""
+        batch_size, seq_len = sequence.shape
+        device = sequence.device
+        
+        if not self.has_mappings:
+            return torch.zeros_like(sequence)
+
+        sequence = torch.clamp(sequence, 0, self.vocab_size - 1)
+        
+        genre_sequence = torch.zeros_like(sequence)
+        flat_sequence = sequence.flatten()
+        flat_genre = torch.zeros_like(flat_sequence)
+        
+        token_mask = torch.isin(flat_sequence, self.token_ids)
+        
+        if token_mask.any():
+            valid_tokens = flat_sequence[token_mask]
+
+            with self._cache_lock:
+                cache_key = (device, len(self.token_ids))
+                if cache_key not in self._genre_cache:
+                    sorted_indices = torch.argsort(self.token_ids)
+                    self._genre_cache[cache_key] = {
+                        'sorted_tokens': self.token_ids[sorted_indices],
+                        'sorted_genres': self.genre_ids[sorted_indices]
+                    }
+                
+                cached_data = self._genre_cache[cache_key]
+            
+            indices = torch.searchsorted(cached_data['sorted_tokens'], valid_tokens)
+            indices = torch.clamp(indices, 0, len(cached_data['sorted_tokens']) - 1)
+            exact_matches = cached_data['sorted_tokens'][indices] == valid_tokens
+            
+            genre_values = torch.where(
+                exact_matches, 
+                cached_data['sorted_genres'][indices], 
+                torch.tensor(0, device=device, dtype=self.genre_ids.dtype)
+            )
+            
+            flat_genre[token_mask] = genre_values
+        
+        return flat_genre.view(batch_size, seq_len)
+    
+    def _get_multi_genre_mapping(self, sequence):
+        """Get multiple genres for each anime in sequence with bounds checking"""
+        batch_size, seq_len = sequence.shape
+        device = sequence.device
+        
+        if not self.has_mappings:
+            return torch.zeros(batch_size, seq_len, self.max_genres_per_anime, device=device, dtype=torch.long)
+
+        sequence = torch.clamp(sequence, 0, self.vocab_size - 1)
+
+        genre_sequences = torch.zeros(batch_size, seq_len, self.max_genres_per_anime, device=device, dtype=torch.long)
+        
+        flat_sequence = sequence.flatten()
+        flat_genres = torch.zeros(len(flat_sequence), self.max_genres_per_anime, device=device, dtype=torch.long)
+        
+        token_mask = torch.isin(flat_sequence, self.token_ids)
+        
+        if token_mask.any():
+            valid_tokens = flat_sequence[token_mask]
+
+            with self._cache_lock:
+                cache_key = (device, len(self.token_ids), 'multi')
+                if cache_key not in self._genre_cache:
+                    sorted_indices = torch.argsort(self.token_ids)
+                    self._genre_cache[cache_key] = {
+                        'sorted_tokens': self.token_ids[sorted_indices],
+                        'sorted_genres': self.genre_ids[sorted_indices]  # Shape: (num_tokens, max_genres_per_anime)
+                    }
+                
+                cached_data = self._genre_cache[cache_key]
+            
+            indices = torch.searchsorted(cached_data['sorted_tokens'], valid_tokens)
+            indices = torch.clamp(indices, 0, len(cached_data['sorted_tokens']) - 1)
+            exact_matches = cached_data['sorted_tokens'][indices] == valid_tokens
+
+            genre_values = cached_data['sorted_genres'][indices]  # Shape: (num_valid_tokens, max_genres_per_anime)
+
+            valid_mask = token_mask.nonzero(as_tuple=True)[0]
+            exact_valid_mask = valid_mask[exact_matches]
+            
+            flat_genres[exact_valid_mask] = genre_values[exact_matches]
+        
+        return flat_genres.view(batch_size, seq_len, self.max_genres_per_anime)
+    
+    def _aggregate_genre_embeddings(self, genre_embeddings):
+        """Aggregate multiple genre embeddings per anime"""
+        # genre_embeddings shape: (batch_size, seq_len, max_genres_per_anime, embed_size)
+        batch_size, seq_len, max_genres, embed_size = genre_embeddings.shape
+
+        weights = F.softmax(self.genre_aggregation, dim=0)
+        weighted_genres = torch.einsum('bsgd,g->bsd', genre_embeddings, weights)
+        
+        return weighted_genres
+    
+    def forward(self, sequence):
+        """
+        Enhanced forward pass with per-anime genre processing
+        """
+        if sequence.max() >= self.vocab_size:
+            print(f"Warning: Input contains tokens >= vocab_size ({self.vocab_size})")
+            
+        sequence = torch.clamp(sequence, 0, self.vocab_size - 1)
+
+        token_emb = self.token(sequence)
+
+        if self.multi_genre:
+            genre_sequences = self._get_multi_genre_mapping(sequence)  # (batch, seq, max_genres)
+            
+            genre_sequences = torch.clamp(genre_sequences, 0, self.num_genres - 1)
+            
+            genre_embeddings = self.genre_embed(genre_sequences)  # (batch, seq, max_genres, embed_size)
+            
+            aggregated_genre_emb = self._aggregate_genre_embeddings(genre_embeddings)  # (batch, seq, embed_size)
+
+            combined = torch.cat([token_emb, aggregated_genre_emb], dim=-1)
+        else:
+            genre_sequence = self._get_single_genre_mapping(sequence)
+            
+            genre_sequence = torch.clamp(genre_sequence, 0, self.num_genres - 1)
+            
+            genre_emb = self.genre_embed(genre_sequence)
+            combined = torch.cat([token_emb, genre_emb], dim=-1)
+
+        x = self.fusion_layer(combined)
+        
+        return self.dropout(x)
+
+    
+    def clear_cache(self):
+        """Clear internal caches to free GPU memory"""
+        with self._cache_lock:
+            self._genre_cache.clear()
+    
+    @classmethod
+    def clear_global_cache(cls):
+        """Clear global mappings cache"""
+        with cls._cache_lock:
+            cls._mappings_cache = None

models/bert_modules/embedding/position.py

@@ -0,0 +1,16 @@
+import torch.nn as nn
+import torch
+import math
+
+
+class PositionalEmbedding(nn.Module):
+
+    def __init__(self, max_len, d_model):
+        super().__init__()
+
+        # Compute the positional encodings once in log space.
+        self.pe = nn.Embedding(max_len, d_model)
+
+    def forward(self, x):
+        batch_size = x.size(0)
+        return self.pe.weight.unsqueeze(0).repeat(batch_size, 1, 1)

models/bert_modules/embedding/segment.py

@@ -0,0 +1,6 @@
+import torch.nn as nn
+
+
+class SegmentEmbedding(nn.Embedding):
+    def __init__(self, embed_size=512):
+        super().__init__(3, embed_size, padding_idx=0)

models/bert_modules/embedding/token.py

@@ -0,0 +1,6 @@
+import torch.nn as nn
+
+
+class TokenEmbedding(nn.Embedding):
+    def __init__(self, vocab_size, embed_size=512):
+        super().__init__(vocab_size, embed_size, padding_idx=0)

models/bert_modules/transformer.py

@@ -0,0 +1,31 @@
+import torch.nn as nn
+
+from .attention import MultiHeadedAttention
+from .utils import SublayerConnection, PositionwiseFeedForward
+
+
+class TransformerBlock(nn.Module):
+    """
+    Bidirectional Encoder = Transformer (self-attention)
+    Transformer = MultiHead_Attention + Feed_Forward with sublayer connection
+    """
+
+    def __init__(self, hidden, attn_heads, feed_forward_hidden, dropout):
+        """
+        :param hidden: hidden size of transformer
+        :param attn_heads: head sizes of multi-head attention
+        :param feed_forward_hidden: feed_forward_hidden, usually 4*hidden_size
+        :param dropout: dropout rate
+        """
+
+        super().__init__()
+        self.attention = MultiHeadedAttention(h=attn_heads, d_model=hidden, dropout=dropout)
+        self.feed_forward = PositionwiseFeedForward(d_model=hidden, d_ff=feed_forward_hidden, dropout=dropout)
+        self.input_sublayer = SublayerConnection(size=hidden, dropout=dropout)
+        self.output_sublayer = SublayerConnection(size=hidden, dropout=dropout)
+        self.dropout = nn.Dropout(p=dropout)
+
+    def forward(self, x, mask):
+        x = self.input_sublayer(x, lambda _x: self.attention.forward(_x, _x, _x, mask=mask))
+        x = self.output_sublayer(x, self.feed_forward)
+        return self.dropout(x)

models/bert_modules/utils/__init__.py

@@ -0,0 +1,4 @@
+from .feed_forward import PositionwiseFeedForward
+from .layer_norm import LayerNorm
+from .sublayer import SublayerConnection
+from .gelu import GELU

models/bert_modules/utils/feed_forward.py

@@ -0,0 +1,16 @@
+import torch.nn as nn
+from .gelu import GELU
+
+
+class PositionwiseFeedForward(nn.Module):
+    "Implements FFN equation."
+
+    def __init__(self, d_model, d_ff, dropout=0.1):
+        super(PositionwiseFeedForward, self).__init__()
+        self.w_1 = nn.Linear(d_model, d_ff)
+        self.w_2 = nn.Linear(d_ff, d_model)
+        self.dropout = nn.Dropout(dropout)
+        self.activation = GELU()
+
+    def forward(self, x):
+        return self.w_2(self.dropout(self.activation(self.w_1(x))))

models/bert_modules/utils/gelu.py

@@ -0,0 +1,12 @@
+import torch.nn as nn
+import torch
+import math
+
+
+class GELU(nn.Module):
+    """
+    Paper Section 3.4, last paragraph notice that BERT used the GELU instead of RELU
+    """
+
+    def forward(self, x):
+        return 0.5 * x * (1 + torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * torch.pow(x, 3))))

models/bert_modules/utils/layer_norm.py

@@ -0,0 +1,17 @@
+import torch.nn as nn
+import torch
+
+
+class LayerNorm(nn.Module):
+    "Construct a layernorm module (See citation for details)."
+
+    def __init__(self, features, eps=1e-6):
+        super(LayerNorm, self).__init__()
+        self.a_2 = nn.Parameter(torch.ones(features))
+        self.b_2 = nn.Parameter(torch.zeros(features))
+        self.eps = eps
+
+    def forward(self, x):
+        mean = x.mean(-1, keepdim=True)
+        std = x.std(-1, keepdim=True)
+        return self.a_2 * (x - mean) / (std + self.eps) + self.b_2

models/bert_modules/utils/sublayer.py

@@ -0,0 +1,18 @@
+import torch.nn as nn
+from .layer_norm import LayerNorm
+
+
+class SublayerConnection(nn.Module):
+    """
+    A residual connection followed by a layer norm.
+    Note for code simplicity the norm is first as opposed to last.
+    """
+
+    def __init__(self, size, dropout):
+        super(SublayerConnection, self).__init__()
+        self.norm = LayerNorm(size)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x, sublayer):
+        "Apply residual connection to any sublayer with the same size."
+        return x + self.dropout(sublayer(self.norm(x)))

models/dae.py

@@ -0,0 +1,54 @@
+from .base import BaseModel
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class DAEModel(BaseModel):
+    def __init__(self, args):
+        super().__init__(args)
+
+        # Input dropout
+        self.input_dropout = nn.Dropout(p=args.dae_dropout)
+
+        # Construct a list of dimensions for the encoder and the decoder
+        dims = [args.dae_hidden_dim] * 2 * args.dae_num_hidden
+        dims = [args.num_items] + dims + [args.dae_latent_dim]
+
+        # Stack encoders and decoders
+        encoder_modules, decoder_modules = [], []
+        for i in range(len(dims)//2):
+            encoder_modules.append(nn.Linear(dims[2*i], dims[2*i+1]))
+            decoder_modules.append(nn.Linear(dims[-2*i-1], dims[-2*i-2]))
+        self.encoder = nn.ModuleList(encoder_modules)
+        self.decoder = nn.ModuleList(decoder_modules)
+
+        # Initialize weights
+        self.encoder.apply(self.weight_init)
+        self.decoder.apply(self.weight_init)
+
+    def weight_init(self, m):
+        if isinstance(m, nn.Linear):
+            nn.init.kaiming_normal_(m.weight)
+            m.bias.data.normal_(0.0, 0.001)
+
+    @classmethod
+    def code(cls):
+        return 'dae'
+
+    def forward(self, x):
+        x = F.normalize(x)
+        x = self.input_dropout(x)
+        
+        for i, layer in enumerate(self.encoder):
+            x = layer(x)
+            x = torch.tanh(x)
+        
+        for i, layer in enumerate(self.decoder):
+            x = layer(x)
+            if i != len(self.decoder)-1:
+                x = torch.tanh(x)
+
+        return x
+

models/vae.py

@@ -0,0 +1,69 @@
+from .base import BaseModel
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class VAEModel(BaseModel):
+    def __init__(self, args):
+        super().__init__(args)
+        self.latent_dim = args.vae_latent_dim
+
+        # Input dropout
+        self.input_dropout = nn.Dropout(p=args.vae_dropout)
+
+        # Construct a list of dimensions for the encoder and the decoder
+        dims = [args.vae_hidden_dim] * 2 * args.vae_num_hidden
+        dims = [args.num_items] + dims + [args.vae_latent_dim * 2]
+
+        # Stack encoders and decoders
+        encoder_modules, decoder_modules = [], []
+        for i in range(len(dims)//2):
+            encoder_modules.append(nn.Linear(dims[2*i], dims[2*i+1]))
+            if i == 0:
+                decoder_modules.append(nn.Linear(dims[-1]//2, dims[-2]))
+            else:
+                decoder_modules.append(nn.Linear(dims[-2*i-1], dims[-2*i-2]))
+        self.encoder = nn.ModuleList(encoder_modules)
+        self.decoder = nn.ModuleList(decoder_modules)
+
+        # Initialize weights
+        self.encoder.apply(self.weight_init)
+        self.decoder.apply(self.weight_init)
+
+    def weight_init(self, m):
+        if isinstance(m, nn.Linear):
+            nn.init.kaiming_normal_(m.weight)
+            m.bias.data.zero_()
+
+    @classmethod
+    def code(cls):
+        return 'vae'
+
+    def forward(self, x):
+        x = F.normalize(x)
+        x = self.input_dropout(x)
+        
+        for i, layer in enumerate(self.encoder):
+            x = layer(x)
+            if i != len(self.encoder) - 1:
+                x = torch.tanh(x)
+        
+        mu, logvar = x[:, :self.latent_dim], x[:, self.latent_dim:]
+
+        if self.training:
+            # since log(var) = log(sigma^2) = 2*log(sigma)
+            sigma = torch.exp(0.5 * logvar)
+            eps = torch.randn_like(sigma)
+            x = mu + eps * sigma
+        else:
+            x = mu
+
+        for i, layer in enumerate(self.decoder):
+            x = layer(x)
+            if i != len(self.decoder) - 1:
+                x = torch.tanh(x)
+                
+        return x, mu, logvar
+