update model AI

main.py

@@ -29,17 +29,20 @@ class HeartRateData(BaseModel):
 # ==========================================
 # Cấu trúc class phải y hệt lúc train
 class BioVibeAI(nn.Module):
-    def __init__(self, vocab_size, embed_size=256, hidden_size=512, num_layers=2):
-        super(BioVibeAI, self).__init__()
-        self.embedding = nn.Embedding(vocab_size, embed_size)
-        self.lstm = nn.LSTM(embed_size, hidden_size, num_layers, batch_first=True)
-        self.fc = nn.Linear(hidden_size, vocab_size)
-        
+    def __init__(self, vocab_size, embed=128, hidden=256):
+        super().__init__()
+
+        self.embedding = nn.Embedding(vocab_size, embed)
+        self.lstm = nn.LSTM(embed, hidden, num_layers=2, batch_first=True, dropout=0.3)
+        self.norm = nn.LayerNorm(hidden)
+        self.fc = nn.Linear(hidden, vocab_size)
+
     def forward(self, x):
-        embedded = self.embedding(x)
-        out, _ = self.lstm(embedded)
-        out = self.fc(out[:, -1, :])
-        return out
+        x = self.embedding(x)
+        out, _ = self.lstm(x)
+        out = out[:, -1, :]
+        out = self.norm(out)
+        return self.fc(out)
 
 # Nạp từ điển
 with open('mapping_dict.pkl', 'rb') as f: