Update app.py

app.py

@@ -1,24 +1,34 @@
-# main.py
-from fastapi import FastAPI
-from pydantic import BaseModel
-import onnxruntime as ort
+import gradio as gr
 import numpy as np
+import onnxruntime as ort
+
+# Load the ONNX model
+session = ort.InferenceSession("model.onnx", providers=["CPUExecutionProvider"])
+
+# 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)
 
-app = FastAPI()
+    # Run the model
+    outputs = session.run(None, {
+        "input_ids": input_ids_np,
+        "attention_mask": attention_mask_np
+    })
 
-session = ort.InferenceSession("model.onnx")
+    # Return raw outputs or post-process as needed
+    return outputs
 
-class ModelInput(BaseModel):
-    input_ids: list[int]
-    attention_mask: list[int]
+# Expose API endpoint
+demo = gr.Interface(
+    fn=predict,
+    inputs=[
+        gr.JSON(label="input_ids"),
+        gr.JSON(label="attention_mask")
+    ],
+    outputs="json",
+    allow_flagging="never"
+)
 
-@app.post("/predict")
-def predict(data: ModelInput):
-    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)
-    inputs = {
-        "input_ids": input_ids,
-        "attention_mask": attention_mask,
-    }
-    outputs = session.run(None, inputs)
-    return {"output": outputs}
+app = gr.mount_gradio_app(app=None, blocks=demo, path="/")