Merge train.py with generate.py

app.py

@@ -8,8 +8,75 @@ import numpy as np
 import os
 import time
 
-# Import the model architecture from train.py
-from train import CVAE, TextEncoder, LATENT_DIM, HIDDEN_DIM
+LATENT_DIM = 128
+HIDDEN_DIM = 256
+
+
+# Text encoder
+class TextEncoder(nn.Module):
+    def __init__(self, hidden_size, output_size):
+        super(TextEncoder, self).__init__()
+        self.bert = BertModel.from_pretrained('bert-base-uncased')
+        self.fc = nn.Linear(self.bert.config.hidden_size, output_size)
+
+    def forward(self, input_ids, attention_mask):
+        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
+        return self.fc(outputs.last_hidden_state[:, 0, :])
+
+# CVAE model
+class CVAE(nn.Module):
+    def __init__(self, text_encoder):
+        super(CVAE, self).__init__()
+        self.text_encoder = text_encoder
+
+        # Encoder
+        self.encoder = nn.Sequential(
+            nn.Conv2d(4, 32, 3, stride=1, padding=1),
+            nn.ReLU(),
+            nn.Conv2d(32, 64, 3, stride=2, padding=1),
+            nn.ReLU(),
+            nn.Conv2d(64, 128, 3, stride=2, padding=1),
+            nn.ReLU(),
+            nn.Flatten(),
+            nn.Linear(128 * 4 * 4, HIDDEN_DIM)
+        )
+
+        self.fc_mu = nn.Linear(HIDDEN_DIM + HIDDEN_DIM, LATENT_DIM)
+        self.fc_logvar = nn.Linear(HIDDEN_DIM + HIDDEN_DIM, LATENT_DIM)
+
+        # Decoder
+        self.decoder_input = nn.Linear(LATENT_DIM + HIDDEN_DIM, 128 * 4 * 4)
+        self.decoder = nn.Sequential(
+            nn.ConvTranspose2d(128, 64, 3, stride=2, padding=1, output_padding=1),
+            nn.ReLU(),
+            nn.ConvTranspose2d(64, 32, 3, stride=2, padding=1, output_padding=1),
+            nn.ReLU(),
+            nn.Conv2d(32, 4, 3, stride=1, padding=1),
+            nn.Tanh()
+        )
+
+    def encode(self, x, c):
+        x = self.encoder(x)
+        x = torch.cat([x, c], dim=1)
+        mu = self.fc_mu(x)
+        logvar = self.fc_logvar(x)
+        return mu, logvar
+
+    def decode(self, z, c):
+        z = torch.cat([z, c], dim=1)
+        x = self.decoder_input(z)
+        x = x.view(-1, 128, 4, 4)
+        return self.decoder(x)
+
+    def reparameterize(self, mu, logvar):
+        std = torch.exp(0.5 * logvar)
+        eps = torch.randn_like(std)
+        return mu + eps * std
+
+    def forward(self, x, c):
+        mu, logvar = self.encode(x, c)
+        z = self.reparameterize(mu, logvar)
+        return self.decode(z, c), mu, logvar
 
 # Initialize the BERT tokenizer
 tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')