Update app.py

app.py

@@ -1,28 +1,22 @@
-from transformers import AutoTokenizer
+from fastapi import FastAPI
 from onnxruntime import InferenceSession
 import numpy as np
-import json
-from fastapi import FastAPI
 
 app = FastAPI()
 
-# Initialize components
-tokenizer = AutoTokenizer.from_pretrained(
-    "Xenova/multi-qa-mpnet-base-dot-v1",
-    use_fast=False  # Avoids framework dependencies
-)
+# Load ONNX model only
 session = InferenceSession("model.onnx")
 
-def cosine_similarity(a, b):
-    return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))
-
 @app.post("/predict")
-async def predict(query: str):
-    # Tokenize
-    inputs = tokenizer(query, return_tensors="np")
-    inputs = {k: v.astype(np.int64) for k, v in inputs.items()}
+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)
     
-    # Get embedding
-    embedding = session.run(None, inputs)[0][0]
+    # Run model
+    outputs = session.run(None, {
+        "input_ids": input_ids,
+        "attention_mask": attention_mask
+    })
     
-    return {"embedding": embedding.tolist()}
+    return {"embedding": outputs[0].tolist()}