feat: autoencoder + updated trainers + cleaned up process to allow using autoencoder

src/models/__init__.py

@@ -1,10 +1,8 @@
 from .Model import Model
 from .autoencoder import AutoEncoder
 from .cnn import CNNPredictor
-from .transformer import ByteTransformer
 
 model_called: dict[str, type[Model]] = {
     'cnn': CNNPredictor,
-    'transformer': ByteTransformer,
     'autoencoder': AutoEncoder
 }

src/models/autoencoder/autoencoder.py

@@ -47,9 +47,15 @@ class AutoEncoder(Model):
         self.decoder = Decoder(latent_dim, channel_count, input_size)
 
     def encode(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        x: torch.Tensor of floats
+        """
         return self.encoder(x)
 
     def decode(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        x: torch.Tensor of floats
+        """
         return self.decoder(x)
 
     def forward(self, x: torch.LongTensor) -> torch.Tensor:

src/models/transformer/__init__.py

@@ -1 +0,0 @@
-from .transformer import ByteTransformer

src/models/transformer/transformer.py

@@ -1,70 +0,0 @@
-from typing import Optional
-
-import torch.nn as nn
-from torch import Tensor, arange
-
-from src.models import Model
-
-
-class LearnedPositionalEncoding(Model):
-    def __init__(self, max_len, d_model):
-        super().__init__()
-        self.pos_emb = nn.Embedding(max_len, d_model)
-
-    def forward(self, x):
-        # x: [seq, batch, d_model]
-        seq_len = x.size(0)
-        positions = arange(seq_len, device=x.device).unsqueeze(1)  # [seq, 1]
-        return x + self.pos_emb(positions)  # broadcast over batch
-
-class ByteTransformer(nn.Module):
-    def __init__(
-            self,
-            d_model=512,
-            nhead=8,
-            num_encoder_layers=6,
-            num_decoder_layers=6,
-            dim_feedforward=2048,
-            dropout=0.1,
-            activation="relu",
-            layer_norm_eps=1e-05,
-            max_len=128
-    ):
-        super().__init__()
-        self.src_embedding = nn.Embedding(256, d_model)
-        self.tgt_embedding = nn.Embedding(256, d_model)
-
-        self.src_pos = LearnedPositionalEncoding(max_len, d_model)
-        self.tgt_pos = LearnedPositionalEncoding(max_len, d_model)
-
-        self.transformer = nn.Transformer(
-            d_model=d_model,
-            nhead=nhead,
-            num_encoder_layers=num_encoder_layers,
-            num_decoder_layers=num_decoder_layers,
-            dim_feedforward=dim_feedforward,
-            dropout=dropout,
-            activation=activation,
-            layer_norm_eps=layer_norm_eps,
-            batch_first=False,
-            norm_first=False,
-            device=None,
-            dtype=None,
-        )
-
-        self.output_proj = nn.Linear(d_model, 256)
-
-        self.loss_function = nn.CrossEntropyLoss()
-
-    def forward(
-            self,
-            src: Tensor,
-            tgt: Tensor,
-    ) -> Tensor:
-        src_embeds = self.src_embedding(src)
-        tgt_embeds = self.tgt_embedding(tgt)
-
-        src_pos = self.src_pos(src_embeds)
-        tgt_pos = self.tgt_pos(tgt_embeds)
-
-        return self.output_proj(self.transformer(src_pos, tgt_pos))