Update server.py: v3-fixed model priority, generalization API + UI

server.py

@@ -22,14 +22,16 @@ STATIC_DIR = os.path.join(BASE_DIR, "static")
 os.makedirs(STATIC_DIR, exist_ok=True)
 
 HF_REPO  = "Sbhat2026/protfunc-models"
-# Priority order: v3 supplemented > improved base > original baseline
+# Priority order: v3_fixed (ablation best) > v3 supplemented > improved base > original baseline
 HF_FILES = [
+    "protfunc_v3_fixed.pth", "protfunc_v3_fixed_thresholds.json",
     "protfunc_v3.pth", "protfunc_v3_thresholds.json",
     "improved_res.pth", "improved_per_label_thresholds.json",
     "baseline_res.pth", "mlb_public_v1.pkl", "go_annotations_fixed.csv", "go_names.json",
 ]
 OPTIONAL = {
     "go_names.json",
+    "protfunc_v3_fixed.pth", "protfunc_v3_fixed_thresholds.json",
     "protfunc_v3.pth", "protfunc_v3_thresholds.json",
     "improved_res.pth", "improved_per_label_thresholds.json",
 }
@@ -215,6 +217,7 @@ def load_go_map():
 
 def load_thresholds():
     for path in [
+        os.path.join(BASE_DIR, "protfunc_v3_fixed_thresholds.json"),
         os.path.join(BASE_DIR, "improved_per_label_thresholds.json"),
         os.path.join(BASE_DIR, "protfunc_v3_thresholds.json"),
         os.path.join(BASE_DIR, "per_label_thresholds.json"),
@@ -518,10 +521,9 @@ async def lifespan(app: FastAPI):
     import numpy as np
     device = torch.device("cpu")
 
-    # Prefer checkpoints in priority order: improved > v3 > baseline
-    # (v3 was trained with propagated training labels which inflates predictions;
-    #  improved_res.pth has verified Fmax=0.8846 on specific GO annotations)
+    # Prefer checkpoints in priority order: v3_fixed (ablation best, CAFA-correct) > improved > v3 > baseline
     ckpt_candidates = [
+        os.path.join(BASE_DIR, "protfunc_v3_fixed.pth"),
         os.path.join(BASE_DIR, "improved_res.pth"),
         os.path.join(BASE_DIR, "protfunc_v3.pth"),
         os.path.join(BASE_DIR, "baseline_res.pth"),
@@ -591,10 +593,12 @@ async def root():
 @app.get("/api/model/info")
 async def model_info():
     """Return model metadata and configuration."""
+    v3_fixed = os.path.exists(os.path.join(BASE_DIR, "protfunc_v3_fixed.pth"))
     improved = os.path.exists(os.path.join(BASE_DIR, "improved_res.pth"))
-    v3       = os.path.exists(os.path.join(BASE_DIR, "protfunc_v3.pth"))
-    # model name reflects actual loaded model (improved takes priority)
-    if model_uses_supp:
+    # model name reflects actual loaded model (v3_fixed takes highest priority)
+    if v3_fixed and model_uses_supp:
+        name, version, active = "ProtFunc v3-fixed (ablation best, CAFA-correct)", "3.1.0", "protfunc_v3_fixed"
+    elif model_uses_supp:
         name, version, active = "ProtFunc v3 (supplemented + mammal)", "3.0.0", "protfunc_v3"
     elif improved:
         name, version, active = "ProtFunc Enhanced", "2.1.0", "improved"
@@ -618,6 +622,28 @@ async def model_info():
     }
 
 
+@app.get("/api/generalization")
+async def get_generalization():
+    """
+    Return cross-taxon generalization results from eval_generalization.py output.
+    Serves artifacts/generalization_results.json if present, otherwise returns empty.
+    """
+    candidates = [
+        os.path.join(BASE_DIR, "artifacts", "generalization", "generalization_results.json"),
+        os.path.join(BASE_DIR, "generalization_results.json"),
+    ]
+    for path in candidates:
+        if os.path.exists(path):
+            with open(path) as f:
+                data = json.load(f)
+            return {
+                "available": True,
+                "taxa": list(data.keys()),
+                "results": data,
+            }
+    return {"available": False, "taxa": [], "results": {}}
+
+
 @app.get("/api/structure")
 async def get_structure(uniprot_id: str):
     """Look up AlphaFold structure data for a UniProt accession."""