Upload AutoEncoder

config.json

@@ -6,11 +6,12 @@
     "AutoConfig": "modeling_autoencoder.AutoEncoderConfig",
     "AutoModel": "modeling_autoencoder.AutoEncoder"
   },
+  "bidirectional": true,
   "compression_rate": 0.5,
-  "dropout_rate": 0.1,
+  "dropout_rate": 0.5,
   "input_dim": 784,
   "latent_dim": 32,
-  "layer_types": "rnn",
+  "layer_types": "lstm",
   "model_type": "autoencoder",
   "num_layers": 4,
   "torch_dtype": "float32",

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:0843a9219f8207ea933b09d69c8fd229ed196f2427634d4a250e313c2e7cb790
-size 7348552
+oid sha256:61c6f8af940de39412307f20af9dedc7df9e669e2a516638ee565e44d8cb5118
+size 47379808

modeling_autoencoder.py

@@ -1,15 +1,19 @@
 
 from torch import nn
 
-#from autoencoder_model.configuration_autoencoder import AutoEncoderConfig
+from transformers import PretrainedConfig, PreTrainedModel, AutoConfig, AutoModel
 
+# from huggingface_hub import notebook_login
 
+# notebook_login()
 
+# AutoEncoderConfig.register_for_auto_class()
+# AutoEncoder.register_for_auto_class("AutoModel")
 
-from transformers import PretrainedConfig, PreTrainedModel
-
-from transformers import PretrainedConfig
+# AutoConfig.register("autoencoder", AutoEncoderConfig)
+# AutoModel.register(AutoEncoderConfig, AutoModel)
 
+# autoencoder.push_to_hub("autoencoder")
 class AutoEncoderConfig(PretrainedConfig):
     model_type = "autoencoder"
 
@@ -21,6 +25,7 @@ class AutoEncoderConfig(PretrainedConfig):
         dropout_rate=None, 
         num_layers=None, 
         compression_rate=None, 
+        bidirectional=False,
         **kwargs
     ):
         super().__init__(**kwargs)
@@ -30,8 +35,9 @@ class AutoEncoderConfig(PretrainedConfig):
         self.dropout_rate = dropout_rate
         self.num_layers = num_layers
         self.compression_rate = compression_rate
+        self.bidirectional = bidirectional
 
-def create_layers(model_section, layer_types, input_dim, latent_dim, num_layers, dropout_rate, compression_rate):
+def create_layers(model_section, layer_types, input_dim, latent_dim, num_layers, dropout_rate, compression_rate, bidirectional):
 
     layers = []
     current_dim = input_dim
@@ -52,18 +58,18 @@ def create_layers(model_section, layer_types, input_dim, latent_dim, num_layers,
         input_diamensions.reverse()
         output_diamensions.reverse()
 
+        if bidirectional & (layer_types in ['lstm', 'rnn', 'gru']):
+            output_diamensions = [2*value for value in output_diamensions]
+
     for idx, (input_dim, output_dim) in enumerate(zip(input_diamensions, output_diamensions)):
         if layer_types == 'linear':
             layers.append(nn.Linear(input_dim, output_dim))
         elif layer_types == 'lstm':
-            # Assuming we are using LSTMs in a way that returns a sequence output
-            layers.append(nn.LSTM(input_dim, output_dim, batch_first=True))
+            layers.append(nn.LSTM(input_dim, output_dim // (2 if bidirectional else 1), batch_first=True, bidirectional=bidirectional))
         elif layer_types == 'rnn':
-            # Assuming we are using LSTMs in a way that returns a sequence output
-            layers.append(nn.RNN(input_dim, output_dim, batch_first=True))
+            layers.append(nn.RNN(input_dim, output_dim // (2 if bidirectional else 1), batch_first=True, bidirectional=bidirectional))
         elif layer_types == 'gru':
-            # Assuming we are using LSTMs in a way that returns a sequence output
-            layers.append(nn.GRU(input_dim, output_dim, batch_first=True))
+            layers.append(nn.GRU(input_dim, output_dim // (2 if bidirectional else 1), batch_first=True, bidirectional=bidirectional))
         if (idx != num_layers - 1) & (dropout_rate != None):
             layers.append(nn.Dropout(dropout_rate))
     return nn.Sequential(*layers)
@@ -76,19 +82,40 @@ class AutoEncoder(PreTrainedModel):
         
         self.encoder = create_layers("encoder",
             config.layer_types, config.input_dim, config.latent_dim, 
-            config.num_layers, config.dropout_rate, config.compression_rate
+            config.num_layers, config.dropout_rate, config.compression_rate,
+            config.bidirectional,
         )
         # Assuming symmetry between encoder and decoder
         self.decoder = create_layers("decoder",
             config.layer_types, config.input_dim, config.latent_dim, 
-            config.num_layers, config.dropout_rate, config.compression_rate
+            config.num_layers, config.dropout_rate, config.compression_rate,
+            config.bidirectional,
         )
 
     def forward(self, x):
-        # Handle LSTM differently since it outputs (output, (h_n, c_n))
-        if config.layer_types == ['lstm', 'rnn', 'gru']:
-            x, _ = self.encoder(x)
-            x, _ = self.decoder(x)
+        if self.config.layer_types in ['lstm', 'rnn', 'gru']:
+            for layer in self.encoder:
+                print(layer)
+                if isinstance(layer, nn.LSTM):
+                    x, (h_n, c_n)= layer(x)
+
+                elif isinstance(layer, nn.RNN):
+                    x, h_o = layer(x)
+                elif isinstance(layer, nn.GRU):
+                    x, h_o = layer(x)
+                else:
+                    x = layer(x)
+                    
+            for layer in self.decoder:
+                if isinstance(layer, nn.LSTM):
+                    x, (h_n, c_n) = layer(x)
+                elif isinstance(layer, nn.RNN):
+                    x, h_o = layer(x)
+                elif isinstance(layer, nn.GRU):
+                    x, h_o = layer(x)
+                else:
+                    x = layer(x)
+                    
         else:
             x = self.encoder(x)
             x = self.decoder(x)