Update app.py

app.py

@@ -1,22 +1,48 @@
 from fastapi import FastAPI
+from fastapi.middleware.cors import CORSMiddleware
 from onnxruntime import InferenceSession
 import numpy as np
+import os
 
-app = FastAPI()
+app = FastAPI(title="ONNX Model API")
 
-# Load ONNX model only
-# session = InferenceSession("model.onnx")
+# CORS configuration
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# Load ONNX model
+model_path = os.path.join(os.getcwd(), "model.onnx")
+session = InferenceSession(model_path)
+
+@app.get("/")
+def health_check():
+    return {"status": "healthy", "message": "ONNX model is ready"}
 
 @app.post("/predict")
 async def predict(inputs: dict):
-    # Expect pre-tokenized input from client
-    ##input_ids = np.array(inputs["input_ids"], dtype=np.int64)
-    #attention_mask = np.array(inputs["attention_mask"], dtype=np.int64)
+    """Expects {'input_ids': [], 'attention_mask': []}"""
+    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)
+        
+        outputs = session.run(
+            None,
+            {
+                "input_ids": input_ids,
+                "attention_mask": attention_mask
+            }
+        )
+        
+        return {"embedding": outputs[0].tolist()}
     
-    # Run model
-    #outputs = session.run(None, {
-    #    "input_ids": input_ids,
-    #    "attention_mask": attention_mask
-    #})
-    return "Status ok"
-    #return {"embedding": outputs[0].tolist()}
+    except Exception as e:
+        return {"error": str(e)}
+
+# Required for Hugging Face Spaces
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)