update main

main.py

@@ -50,7 +50,7 @@ class TextInput(BaseModel):
     text: str
 
 class BertBiLSTMClassifier(nn.Module):
-    def __init__(self, bert_model_name, num_emotion_classes, binary_cols, lstm_hidden_size=256, dropout=0.3):
+    def __init__(self, bert_model_name, num_emotion_classes, binary_cols, lstm_hidden_size=128, dropout=0.4):
         super().__init__()
         self.bert = AutoModel.from_pretrained(bert_model_name)
         self.lstm = nn.LSTM(
@@ -93,7 +93,7 @@ model = BertBiLSTMClassifier(
     bert_model_name="vinai/phobert-base",
     num_emotion_classes=3,
     binary_cols=binary_cols,
-    lstm_hidden_size=256
+    lstm_hidden_size=128
 ).to(device)
 
 # Load model state dict
@@ -101,11 +101,11 @@ model.load_state_dict(model_state_dict)
 model.eval()
 
 threshold_dict = {
-    'sản phẩm': 0.6,
-    'giá cả': 0.4,
-    'vận chuyển': 0.45,
-    'thái độ và dịch vụ khách hàng': 0.35,
-    'khác': 0.4
+    'sản phẩm': 0.28,
+    'giá cả': 0.58,
+    'vận chuyển': 0.58,
+    'thái độ và dịch vụ khách hàng': 0.70,
+    'khác': 0.6
 }
 
 def predict(text: str):