Update app.py

in safetensors

app.py

@@ -11,6 +11,7 @@ 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  
 
 # -----------------------------
 # Config
@@ -30,10 +31,10 @@ TARGETS = [
 ]
 
 MODEL_PATHS = {
-    "jit": "pd_model.ts",           # optional
-    "state": "pd_model.pt",         # state_dict or full Module
-    "sx": "scaler_x.pkl",
-    "sy": "scaler_y.pkl",
+    "jit": "mode/pd_model.ts",             
+    "state": "mode/pd_model.safetensors",  
+    "sx": "mode/scaler_x.pkl",             
+    "sy": "mode/scaler_y.pkl",             
 }
 
 # -----------------------------
@@ -84,28 +85,27 @@ class PDTabTransformer(nn.Module):
 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 working dir.")
+        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"])
 
-    # Model (jit preferred)
+    # 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"]):
-        state = torch.load(MODEL_PATHS["state"], map_location="cpu")
-        if isinstance(state, nn.Module):
-            model = state
-        else:
-            model = PDTabTransformer(num_features=len(FEATURES), output_dim=len(TARGETS))
-            model.load_state_dict(state)
+        # ✅ Load safetensors
+        state_dict = load_file(MODEL_PATHS["state"], device="cpu")
+        model = PDTabTransformer(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.pt")
+        raise RuntimeError("Model file not found. Provide pd_model.ts or pd_model.safetensors")
     model.eval()
     return model, scaler_x, scaler_y
 
 MODEL, SCALER_X, SCALER_Y = load_model_and_scalers()
 
+
 # -----------------------------
 # Utilities
 # -----------------------------