from fastapi import FastAPI, HTTPException, Request
from fastapi.encoders import jsonable_encoder
from onnxruntime import InferenceSession
from transformers import AutoTokenizer
import numpy as np
import uvicorn

app = FastAPI()

# Initialize tokenizer
tokenizer = AutoTokenizer.from_pretrained(
    "Xenova/multi-qa-mpnet-base-dot-v1",
    use_fast=True,
    legacy=False
)

# Load ONNX model
session = InferenceSession("model.onnx")

def convert_output(value):
    """Recursively convert numpy types to native Python types"""
    if isinstance(value, (np.generic, np.ndarray)):
        if value.size == 1:
            return float(value.item())  # Convert single values to float
        return value.astype(float).tolist()  # Convert arrays to list
    elif isinstance(value, list):
        return [convert_output(x) for x in value]
    elif isinstance(value, dict):
        return {k: convert_output(v) for k, v in value.items()}
    return value

@app.post("/api/predict")
async def predict(request: Request):
    try:
        data = await request.json()
        text = data.get("text", "")
        
        if not text:
            raise HTTPException(status_code=400, detail="No text provided")
        
        # Tokenize input
        inputs = tokenizer(
            text,
            return_tensors="np",
            padding=False,  # Disable padding
            truncation=False,  # Disable truncation
            add_special_tokens=True  # Ensure CLS/SEP tokens
        )
        
        onnx_inputs = {
            "input_ids": np.array(inputs["input_ids"], dtype=np.int64),
            "attention_mask": np.array(inputs["attention_mask"], dtype=np.int64)
        }
        
        outputs = session.run(None, onnx_inputs)
        print("OUTPUTS",outputs)
        # Prepare response with converted types
        return {
            "embedding": outputs[0][0].astype(float).tolist(),
            "input_ids": inputs["input_ids"][0].tolist(),
            "attention_mask": inputs["attention_mask"][0].tolist()
        }
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)