Update app.py

app.py

@@ -1,10 +1,12 @@
-from fastapi import FastAPI
+from fastapi import FastAPI, Request, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
 from onnxruntime import InferenceSession
 import numpy as np
 import os
+import uvicorn
 
-app = FastAPI(title="ONNX Model API")
+# Initialize FastAPI with docs disabled for Spaces
+app = FastAPI(docs_url=None, redoc_url=None)
 
 # CORS configuration
 app.add_middleware(
@@ -15,20 +17,28 @@ app.add_middleware(
 )
 
 # Load ONNX model
-model_path = os.path.join(os.getcwd(), "model.onnx")
-session = InferenceSession(model_path)
+try:
+    session = InferenceSession("model.onnx")
+    print("Model loaded successfully")
+except Exception as e:
+    print(f"Model loading failed: {str(e)}")
+    raise
 
 @app.get("/")
-def health_check():
-    return {"status": "healthy", "message": "ONNX model is ready"}
+async def health_check():
+    return {"status": "ready", "model": "onnx"}
 
-@app.post("/predict")
-async def predict(inputs: dict):
-    """Expects {'input_ids': [], 'attention_mask': []}"""
+@app.post("/api/predict")
+async def predict(request: Request):
     try:
-        input_ids = np.array(inputs["input_ids"], dtype=np.int64).reshape(1, -1)
-        attention_mask = np.array(inputs["attention_mask"], dtype=np.int64).reshape(1, -1)
+        # 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,
             {
@@ -40,9 +50,13 @@ async def predict(inputs: dict):
         return {"embedding": outputs[0].tolist()}
     
     except Exception as e:
-        return {"error": str(e)}
+        raise HTTPException(status_code=400, detail=str(e))
 
 # Required for Hugging Face Spaces
 if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    uvicorn.run(
+        "app:app",
+        host="0.0.0.0",
+        port=7860,
+        reload=False
+    )