Update app.py

app.py

@@ -1,11 +1,12 @@
-from fastapi import FastAPI, Request, HTTPException
+from fastapi import FastAPI, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
+from transformers import AutoTokenizer
 from onnxruntime import InferenceSession
 import numpy as np
 import os
-import uvicorn
+from typing import Dict
 
-app = FastAPI(title="ONNX Model API")
+app = FastAPI(title="ONNX Model API with Tokenizer")
 
 # CORS configuration
 app.add_middleware(
@@ -15,47 +16,50 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# Load ONNX model
+# Initialize components
+tokenizer = AutoTokenizer.from_pretrained("Xenova/multi-qa-mpnet-base-dot-v1")
 session = InferenceSession("model.onnx")
 
-# Essential for Spaces health checks
-@app.get("/")
-def read_root():
-    return {"status": "ONNX Model API is running"}
+def convert_outputs(outputs):
+    """Ensure all numpy values are converted to Python native types"""
+    if isinstance(outputs, (np.generic, np.ndarray)):
+        return outputs.item() if outputs.ndim == 0 else outputs.tolist()
+    return outputs
 
-# Main prediction endpoint
-@app.post("/predict")
-async def predict(request: Request):
+@app.post("/api/process")
+async def process_text(request: Dict[str, str]):
     try:
-        data = await request.json()
-        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)
+        text = request.get("text", "")
         
-        outputs = session.run(None, {
-            "input_ids": input_ids,
-            "attention_mask": attention_mask
-        })
+        # Tokenize the input text
+        inputs = tokenizer(
+            text,
+            return_tensors="np",
+            padding=True,
+            truncation=True,
+            max_length=32  # Match your model's expected input size
+        )
         
-        result = {
-            "embedding": outputs[0].astype(np.float32).tolist()  # Force float32 conversion
+        # Convert to ONNX-compatible format
+        onnx_inputs = {
+            "input_ids": inputs["input_ids"].astype(np.int64),
+            "attention_mask": inputs["attention_mask"].astype(np.int64)
+        }
+        
+        # Run model inference
+        outputs = session.run(None, onnx_inputs)
+        
+        # Convert all numpy types to native Python types
+        processed_outputs = [convert_outputs(output) for output in outputs]
+        
+        return {
+            "embedding": processed_outputs[0],  # Assuming first output is embeddings
+            "tokens": tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
         }
         
-        return jsonable_encoder(result)
-    
     except Exception as e:
         raise HTTPException(status_code=400, detail=str(e))
 
-# Special endpoint for Spaces compatibility
-@app.post("/api/predict")
-async def spaces_predict(request: Request):
-    return await predict(request)
-
-if __name__ == "__main__":
-    uvicorn.run(
-        app,
-        host="0.0.0.0",
-        port=7860,
-        # Required for Spaces:
-        proxy_headers=True,
-        forwarded_allow_ips="*"
-    )
+@app.get("/health")
+async def health_check():
+    return {"status": "healthy"}