Update app.py

merge the csv preditions in response of an batch request

app.py

@@ -1,8 +1,7 @@
-# main.py
 import io
 import os
-from typing import List, Optional
-import joblib
+from typing import List
+import pickle
 import numpy as np
 import pandas as pd
 import torch
@@ -10,12 +9,12 @@ import torch.nn as nn
 from fastapi import FastAPI, File, UploadFile, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
-from safetensors.torch import load_file  
+from safetensors.torch import load_file
 
 # -----------------------------
 # Config
 # -----------------------------
-FEATURES = [
+FEATUREacieS = [
     "ESS_TOTAL", "MCATOT", "GDS_TOTAL",
     "MCAALTTM", "MCACUBE", "MCASER7", "MCAABSTR",
     "GDSSATIS", "GDSHAPPY", "GDSENRGY",
@@ -30,10 +29,9 @@ TARGETS = [
 ]
 
 MODEL_PATHS = {
-    "jit": "mode/pd_model.ts",             
-    "state": "mode/pd_model.safetensors",  
-    "sx": "model/scaler_x.pkl",             
-    "sy": "model/scaler_y.pkl",             
+    "state": "model/flake_transformer.safetensors",
+    "sx": "model/scaler_x.pkl",
+    "sy": "model/scaler_y.pkl",
 }
 
 # -----------------------------
@@ -43,21 +41,21 @@ app = FastAPI(title="PD Biomarker Predictor", version="1.0.0")
 
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],  # tighten for prod
+    allow_origins=["*"],  # Tighten for production
     allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
 )
 
 # -----------------------------
-# Model definition (for state_dict load)
+# Model definition
 # -----------------------------
-class flake_Transformer(nn.Module):
+class flakeParkinsonTransformer(nn.Module):
     def __init__(self, num_features: int, output_dim: int, embed_dim=64, nhead=4, num_layers=3):
         super().__init__()
         self.feature_embeds = nn.ModuleList([nn.Linear(1, embed_dim) for _ in range(num_features)])
         encoder_layer = nn.TransformerEncoderLayer(
-            d_model=embed_dim, nhead=nhead, dim_feedforward=256, dropout=0.1, batch_first=False
+            d_model=embed_dim, nhead=nhead, dim_feedforward=256, dropout=0.1
         )
         self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
         self.pool = nn.AdaptiveAvgPool1d(1)
@@ -70,46 +68,51 @@ class flake_Transformer(nn.Module):
 
     def forward(self, x):
         embeds = [self.feature_embeds[i](x[:, i].unsqueeze(1)) for i in range(x.shape[1])]
-        x = torch.stack(embeds, dim=1)                # [B, F, E]
-        x = x.permute(1, 0, 2)                        # [F, B, E]
-        x = self.transformer(x)                       # [F, B, E]
-        x = x.permute(1, 2, 0)                        # [B, E, F]
-        x = self.pool(x).squeeze(2)                   # [B, E]
-        x = self.fc(x)                                # [B, 4]
+        x = torch.stack(embeds, dim=1)
+        x = x.permute(1, 0, 2)
+        x = self.transformer(x)
+        x = x.permute(1, 2, 0)
+        x = self.pool(x).squeeze(2)
+        x = self.fc(x)
         return x
 
 # -----------------------------
 # Load model and scalers
 # -----------------------------
 def load_model_and_scalers():
-    # Scalers
-    if not (os.path.exists(MODEL_PATHS["sx"]) and os.path.exists(MODEL_PATHS["sy"])):
-        raise RuntimeError("Missing scalers. Expected scaler_x.pkl and scaler_y.pkl in mode/ directory.")
-    scaler_x = joblib.load(MODEL_PATHS["sx"])
-    scaler_y = joblib.load(MODEL_PATHS["sy"])
+    # Load scalers using pickle
+    try:
+        with open(MODEL_PATHS["sx"], 'rb') as f:
+            scaler_x = pickle.load(f)
+        print("scaler_x loaded successfully.")
+    except Exception as e:
+        raise RuntimeError(f"Error loading scaler_x: {e}")
 
-    # Model
-    model = None
-    if os.path.exists(MODEL_PATHS["jit"]):
-        model = torch.jit.load(MODEL_PATHS["jit"], map_location="cpu")
-    elif os.path.exists(MODEL_PATHS["state"]):
-        #  Load safetensors
+    try:
+        with open(MODEL_PATHS["sy"], 'rb') as f:
+            scaler_y = pickle.load(f)
+        print("scaler_y loaded successfully.")
+    except Exception as e:
+        raise RuntimeError(f"Error loading scaler_y: {e}")
+
+    # Load model using safetensors
+    try:
         state_dict = load_file(MODEL_PATHS["state"], device="cpu")
-        model = PDTabTransformer(num_features=len(FEATURES), output_dim=len(TARGETS))
+        model = flakeParkinsonTransformer(num_features=len(FEATURES), output_dim=len(TARGETS))
         model.load_state_dict(state_dict)
-    else:
-        raise RuntimeError("Model file not found. Provide pd_model.ts or pd_model.safetensors")
-    model.eval()
+        model.eval()
+        print("Model loaded and set to evaluation mode.")
+    except Exception as e:
+        raise RuntimeError(f"Error loading model state dictionary: {e}")
+
     return model, scaler_x, scaler_y
 
 MODEL, SCALER_X, SCALER_Y = load_model_and_scalers()
 
-
 # -----------------------------
 # Utilities
 # -----------------------------
 def try_compute_ess_total(df: pd.DataFrame) -> pd.DataFrame:
-    # If ESS_TOTAL missing, sum ESS1..ESS8 when available
     if "ESS_TOTAL" not in df.columns:
         ess_cols = [f"ESS{i}" for i in range(1, 9) if f"ESS{i}" in df.columns]
         if len(ess_cols) == 8:
@@ -117,11 +120,10 @@ def try_compute_ess_total(df: pd.DataFrame) -> pd.DataFrame:
     return df
 
 def try_compute_gds_total(df: pd.DataFrame) -> pd.DataFrame:
-    # If GDS_TOTAL missing, sum standard 15 items when available
     gds_items = [
-        "GDSSATIS","GDSDROPD","GDSEMPTY","GDSBORED","GDSGSPIR",
-        "GDSAFRAD","GDSHAPPY","GDSHLPLS","GDSHOME","GDSMEMRY",
-        "GDSALIVE","GDSWRTLS","GDSENRGY","GDSHOPLS","GDSBETER"
+        "GDSSATIS", "GDSDROPD", "GDSEMPTY", "GDSBORED", "GDSGSPIR",
+        "GDSAFRAD", "GDSHAPPY", "GDSHLPLS", "GDSHOME", "GDSMEMRY",
+        "GDSALIVE", "GDSWRTLS", "GDSENRGY", "GDSHOPLS", "GDSBETER"
     ]
     if "GDS_TOTAL" not in df.columns:
         present = [c for c in gds_items if c in df.columns]
@@ -139,11 +141,9 @@ def standardize_and_predict(df_features: pd.DataFrame):
     return preds
 
 def merge_four_frames(ess: pd.DataFrame, moca: pd.DataFrame, gds: pd.DataFrame, dat: pd.DataFrame) -> pd.DataFrame:
-    # compute totals if needed
     ess = try_compute_ess_total(ess)
     gds = try_compute_gds_total(gds)
 
-    # basic checks
     for df, name in [(ess, "ESS"), (moca, "MoCA"), (gds, "GDS"), (dat, "DaTSCAN")]:
         if not all(col in df.columns for col in ["PATNO", "EVENT_ID"]):
             raise HTTPException(status_code=400, detail=f"{name} CSV missing PATNO or EVENT_ID columns")
@@ -152,7 +152,6 @@ def merge_four_frames(ess: pd.DataFrame, moca: pd.DataFrame, gds: pd.DataFrame,
     df = df.merge(gds, on=["PATNO", "EVENT_ID"], how="inner")
     df = df.merge(dat, on=["PATNO", "EVENT_ID"], how="inner")
 
-    # validate features present / inferable
     missing = [f for f in FEATURES if f not in df.columns]
     if missing:
         raise HTTPException(status_code=400, detail=f"Merged CSVs missing required features: {missing}")
@@ -160,9 +159,8 @@ def merge_four_frames(ess: pd.DataFrame, moca: pd.DataFrame, gds: pd.DataFrame,
     return df
 
 def detect_file_kind(name: str) -> str:
-    """Rudimentary detector to map filename to dataset kind."""
     l = name.lower()
-    if "datscan" in l or "dat" in l and "scan" in l:
+    if "datscan" in l or ("dat" in l and "scan" in l):
         return "datscan"
     if "moca" in l:
         return "moca"
@@ -213,7 +211,6 @@ async def predict_files(files: List[UploadFile] = File(...)):
     if len(files) < 4:
         raise HTTPException(status_code=400, detail="Please upload four CSV files: ESS, MoCA, GDS, DaTSCAN.")
 
-    # Read all
     buckets = {"ess": None, "moca": None, "gds": None, "datscan": None}
     fallback = []
     for f in files:
@@ -229,7 +226,6 @@ async def predict_files(files: List[UploadFile] = File(...)):
         else:
             fallback.append((kind, df))
 
-    # If detection failed, try to auto-assign remaining by column heuristics
     if buckets["ess"] is None:
         candidates = [df for kind, df in fallback if "ESS1" in df.columns or "ESS_TOTAL" in df.columns]
         if candidates:
@@ -251,20 +247,25 @@ async def predict_files(files: List[UploadFile] = File(...)):
         raise HTTPException(status_code=400, detail="Could not identify all four CSVs (ESS, MoCA, GDS, DaTSCAN) by filename/columns.")
 
     merged = merge_four_frames(buckets["ess"], buckets["moca"], buckets["gds"], buckets["datscan"])
-
-    # Predict for all rows; return the first plus summary
     preds = standardize_and_predict(merged)
-    first = preds[0].tolist()
+
+    # Create a list of predictions for each patient
+    results = []
+    for idx, pred in enumerate(preds):
+        result = {
+            "PATNO": int(merged.iloc[idx]["PATNO"]),
+            "EVENT_ID": str(merged.iloc[idx]["EVENT_ID"]),
+            "predicted_biomarkers": {
+                TARGETS[0]: float(pred[0]),
+                TARGETS[1]: float(pred[1]),
+                TARGETS[2]: float(pred[2]),
+                TARGETS[3]: float(pred[3]),
+            }
+        }
+        results.append(result)
 
     return {
-        "predicted_biomarkers": {
-            TARGETS[0]: first[0],
-            TARGETS[1]: first[1],
-            TARGETS[2]: first[2],
-            TARGETS[3]: first[3],
-        },
+        "predictions": results,
         "source": "files",
-        "merged_rows": int(merged.shape[0]),
-        "note": "Multiple rows detected; the first row's prediction is returned. Consider filtering by PATNO/EVENT_ID for per-subject results."
-    }
-
+        "merged_rows": int(merged.shape[0])
+    }