From 6e591bb470106ec7b315fb0f4d6870bbcea56b22 Mon Sep 17 00:00:00 2001
From: RobinMeersman <echteenrobin@gmail.com>
Date: Sat, 13 Dec 2025 15:19:58 +0100
Subject: [PATCH] backup

---
 src/models/__init__.py                      |  4 +++-
 src/models/autoencoder/__init__.py          |  1 +
 src/models/{ => autoencoder}/autoencoder.py | 16 +++++++++-------
 3 files changed, 13 insertions(+), 8 deletions(-)
 create mode 100644 src/models/autoencoder/__init__.py
 rename src/models/{ => autoencoder}/autoencoder.py (76%)

diff --git a/src/models/__init__.py b/src/models/__init__.py
index e329dbc..dfdc5de 100644
--- a/src/models/__init__.py
+++ b/src/models/__init__.py
@@ -1,8 +1,10 @@
 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
+    'transformer': ByteTransformer,
+    'autoencoder': AutoEncoder
 }
diff --git a/src/models/autoencoder/__init__.py b/src/models/autoencoder/__init__.py
new file mode 100644
index 0000000..a4eef00
--- /dev/null
+++ b/src/models/autoencoder/__init__.py
@@ -0,0 +1 @@
+from .autoencoder import AutoEncoder
\ No newline at end of file
diff --git a/src/models/autoencoder.py b/src/models/autoencoder/autoencoder.py
similarity index 76%
rename from src/models/autoencoder.py
rename to src/models/autoencoder/autoencoder.py
index 76335b6..1daf116 100644
--- a/src/models/autoencoder.py
+++ b/src/models/autoencoder/autoencoder.py
@@ -1,6 +1,8 @@
 import torch
 import torch.nn as nn
 
+from src.models import Model
+
 
 class Encoder(nn.Module):
     def __init__(self, input_size, hidden_size, latent_dim):
@@ -23,21 +25,21 @@ class Decoder(nn.Module):
     def __init__(self, input_size, hidden_size, output_size):
         super(Decoder, self).__init__()
         super._decoder = nn.Sequential(*[
-            nn.Linear(input_size),
+            nn.Linear(input_size, 2 * hidden_size),
             nn.ReLU(),
-            nn.BatchNorm1d(input_size),
-            nn.ConvTranspose1d(input_size, 2 * hidden_size, kernel_size=3, stride=2, padding=1, output_padding=1),
             nn.BatchNorm1d(2 * hidden_size),
+            nn.ConvTranspose1d(2 * hidden_size, hidden_size, kernel_size=3, stride=2, padding=1, output_padding=1),
+            nn.BatchNorm1d(hidden_size),
             nn.ReLU(),
-            nn.ConvTranspose1d(2 * hidden_size, output_size, kernel_size=3, padding=1),
+            nn.ConvTranspose1d(hidden_size, output_size, kernel_size=3, padding=1),
         ])
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        pass
+        return self._decoder(x)
 
-class AutoEncoder(nn.Module):
+class AutoEncoder(Model):
     def __init__(self, input_size, hidden_size, latent_dim):
-        super(AutoEncoder, self).__init__()
+        super().__init__(loss_function = nn.CrossEntropyLoss())
 
         self.encoder = Encoder(input_size, hidden_size, latent_dim)
         self.decoder = Decoder(latent_dim, hidden_size, input_size)