Update app.py

app.py

@@ -29,23 +29,29 @@ token2id["<EOS>"] = 1
 id2token = {i: t for t, i in token2id.items()}
 VOCAB_SIZE = len(token2id)
 
-# --- Predictor Model Architecture (from antioxidant_predictor_5.py) ---
+# --- Predictor Model Architecture (Corrected to match saved weights) ---
 class AntioxidantPredictor(nn.Module):
-    # This class definition should be an exact copy from your project
     def __init__(self, input_dim, transformer_layers=3, transformer_heads=4, transformer_dropout=0.1):
         super(AntioxidantPredictor, self).__init__()
-        self.input_dim = input_dim
+        # 根据错误日志和您的训练脚本，我们知道输入维度是固定的
+        # 并且模型内部处理 ProtT5 和传统特征的分离
         self.t5_dim = 1024
-        self.hand_crafted_dim = self.input_dim - self.t5_dim
+        self.hand_crafted_dim = input_dim - self.t5_dim
         
+        # 定义 Transformer Encoder
         encoder_layer = nn.TransformerEncoderLayer(
-            d_model=self.t5_dim, nhead=transformer_heads, 
-            dropout=transformer_dropout, batch_first=True
+            d_model=self.t5_dim, 
+            nhead=transformer_heads, 
+            dropout=transformer_dropout, 
+            batch_first=True
         )
         self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=transformer_layers)
         
+        # 定义 MLP
+        # 错误日志表明权重文件没有 fusion_fc 和 classifier，只有一个 mlp
+        # 我们根据 predictor_train.py 的原始结构来重建
         self.mlp = nn.Sequential(
-            nn.Linear(self.input_dim, 512),
+            nn.Linear(input_dim, 512),
             nn.ReLU(),
             nn.Dropout(0.5),
             nn.Linear(512, 256),
@@ -56,18 +62,29 @@ class AntioxidantPredictor(nn.Module):
         self.temperature = nn.Parameter(torch.ones(1))
 
     def forward(self, fused_features):
-        tr_features = fused_features[:, :self.t5_dim]
-        hand_features = fused_features[:, self.t5_dim:]
-        tr_features_unsqueezed = tr_features.unsqueeze(1)
-        transformer_output = self.transformer_encoder(tr_features_unsqueezed)
+        # 这个前向传播逻辑与您的训练脚本 predictor_train.py 更为匹配
+        prot_t5_features = fused_features[:, :self.t5_dim]
+        hand_crafted_features = fused_features[:, self.t5_dim:]
+        
+        # Transformer 只处理 ProtT5 特征
+        prot_t5_features_unsqueezed = prot_t5_features.unsqueeze(1)
+        transformer_output = self.transformer_encoder(prot_t5_features_unsqueezed)
         transformer_output_pooled = transformer_output.mean(dim=1)
-        combined_features = torch.cat((transformer_output_pooled, hand_features), dim=1)
+        
+        # 将处理后的 ProtT5 特征与传统特征拼接
+        combined_features = torch.cat((transformer_output_pooled, hand_crafted_features), dim=1)
+        
+        # 将最终拼接的特征送入 MLP
         logits = self.mlp(combined_features)
+        
         return logits / self.temperature
 
     def get_temperature(self):
         return self.temperature.item()
 
+    def set_temperature(self, temp_value, device):
+        self.temperature = nn.Parameter(torch.tensor([temp_value], device=device), requires_grad=False)
+
 # --- Generator Model Architecture (from generator.py) ---
 class ProtT5Generator(nn.Module):
     # This class definition should be an exact copy from your project