Update app.py

app.py

@@ -1,34 +1,29 @@
-import gradio as gr
+from fastapi import FastAPI
+from fastapi.middleware.cors import CORSMiddleware
 import numpy as np
-import onnxruntime as ort
+from onnxruntime import InferenceSession
+from transformers import AutoTokenizer
+import os
 
-# Load the ONNX model
-session = ort.InferenceSession("model.onnx", providers=["CPUExecutionProvider"])
+app = FastAPI()
 
-# Prediction function
-def predict(input_ids: list[int], attention_mask: list[int]):
-    # Convert to numpy arrays and batch them
-    input_ids_np = np.array([input_ids], dtype=np.int64)
-    attention_mask_np = np.array([attention_mask], dtype=np.int64)
-
-    # Run the model
-    outputs = session.run(None, {
-        "input_ids": input_ids_np,
-        "attention_mask": attention_mask_np
-    })
-
-    # Return raw outputs or post-process as needed
-    return outputs
-
-# Expose API endpoint
-demo = gr.Interface(
-    fn=predict,
-    inputs=[
-        gr.JSON(label="input_ids"),
-        gr.JSON(label="attention_mask")
-    ],
-    outputs="json",
-    allow_flagging="never"
+# CORS setup
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_methods=["*"],
+    allow_headers=["*"],
 )
 
-app = gr.mount_gradio_app(app=None, blocks=demo, path="/")
+# Load model
+session = InferenceSession("model.onnx")
+tokenizer = AutoTokenizer.from_pretrained("Xenova/multi-qa-mpnet-base-dot-v1")
+
+@app.post("/predict")
+async def predict(query: str):
+    inputs = tokenizer(query, return_tensors="np")
+    inputs = {k: v.astype(np.int64) for k, v in inputs.items()}
+    outputs = session.run(None, inputs)
+    embedding = outputs[0][0].tolist()
+    
+    return {"embedding": embedding}