Update utils/inference_ESM2.py

utils/inference_ESM2.py

@@ -56,7 +56,7 @@ def esm2_embed_batch(seqs, tokenizer, model, device, batch_size=8):
 @torch.no_grad()
 def get_final_score(
     epitope,
-    metadata_df,   # 반드시 ['assay','method','state','disease'] string
+    metadata_df,
     tokenizer,
     esm_model,
     device,
@@ -65,37 +65,28 @@ def get_final_score(
     encoder_s2,
     threshold_s1=0.0
 ):
+    FEATURE_ORDER = ['assay', 'method', 'state', 'disease']
+
+    # 🔒 컬럼 순서 강제 (Streamlit rerun 방어)
+    metadata_df = metadata_df[FEATURE_ORDER]
+
     # 1. embedding
-    emb = esm2_embed_batch(
-        [epitope],
-        tokenizer,
-        esm_model,
-        device
-    )
+    emb = esm2_embed_batch([epitope], tokenizer, esm_model, device)
     df_emb = pd.DataFrame(emb)
 
-    # ------------------
-    # Stage 1 (CatBoost)
-    # ------------------
-    X_s1 = pd.concat(
-        [df_emb, metadata_df[['assay','method','state','disease']]],
-        axis=1
-    )
-
+    # Stage 1
+    X_s1 = pd.concat([df_emb, metadata_df], axis=1)
     p1 = model_s1.predict_proba(X_s1)[0, 1]
 
     if p1 < threshold_s1:
         return p1, None, None
 
-    # ------------------
-    # Stage 2 (XGBoost)
-    # ------------------
-    X_cat_s2 = encoder_s2.transform(
-        metadata_df[['assay','method','state','disease']]
-    )
+    # Stage 2
+    X_cat_s2 = encoder_s2.transform(metadata_df)
     X_s2 = np.hstack([df_emb.values, X_cat_s2])
 
     p2 = model_s2.predict_proba(X_s2)[0, 1]
     final = 0.4 * p1 + 0.6 * p2
 
     return p1, p2, final
+