Update app.py

app.py

@@ -5,8 +5,7 @@ import numpy as np
 import os
 import uvicorn
 
-# Initialize FastAPI with docs disabled for Spaces
-app = FastAPI(docs_url=None, redoc_url=None)
+app = FastAPI(title="ONNX Model API")
 
 # CORS configuration
 app.add_middleware(
@@ -17,46 +16,42 @@ app.add_middleware(
 )
 
 # Load ONNX model
-try:
-    session = InferenceSession("model.onnx")
-    print("Model loaded successfully")
-except Exception as e:
-    print(f"Model loading failed: {str(e)}")
-    raise
+session = InferenceSession("model.onnx")
 
+# Essential for Spaces health checks
 @app.get("/")
-async def health_check():
-    return {"status": "ready", "model": "onnx"}
+def read_root():
+    return {"status": "ONNX Model API is running"}
 
-@app.post("/api/predict")
+# Main prediction endpoint
+@app.post("/predict")
 async def predict(request: Request):
     try:
-        # Get JSON input
         data = await request.json()
-        
-        # Convert to numpy arrays with correct shape
         input_ids = np.array(data["input_ids"], dtype=np.int64).reshape(1, -1)
         attention_mask = np.array(data["attention_mask"], dtype=np.int64).reshape(1, -1)
         
-        # Run inference
-        outputs = session.run(
-            None,
-            {
-                "input_ids": input_ids,
-                "attention_mask": attention_mask
-            }
-        )
+        outputs = session.run(None, {
+            "input_ids": input_ids,
+            "attention_mask": attention_mask
+        })
         
         return {"embedding": outputs[0].tolist()}
     
     except Exception as e:
         raise HTTPException(status_code=400, detail=str(e))
 
-# Required for Hugging Face Spaces
+# Special endpoint for Spaces compatibility
+@app.post("/api/predict")
+async def spaces_predict(request: Request):
+    return await predict(request)
+
 if __name__ == "__main__":
     uvicorn.run(
-        "app:app",
+        app,
         host="0.0.0.0",
         port=7860,
-        reload=False
+        # Required for Spaces:
+        proxy_headers=True,
+        forwarded_allow_ips="*"
     )