|
from fastapi import FastAPI, HTTPException |
|
from onnxruntime import InferenceSession |
|
from transformers import AutoTokenizer |
|
import numpy as np |
|
import os |
|
|
|
app = FastAPI() |
|
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained( |
|
"Xenova/multi-qa-mpnet-base-dot-v1", |
|
use_fast=True, |
|
legacy=False |
|
) |
|
|
|
|
|
session = InferenceSession("model.onnx") |
|
|
|
@app.post("/api/predict") |
|
async def predict(text: str): |
|
try: |
|
|
|
inputs = tokenizer( |
|
text, |
|
return_tensors="np", |
|
padding=True, |
|
truncation=True, |
|
max_length=32 |
|
) |
|
|
|
|
|
onnx_inputs = { |
|
"input_ids": inputs["input_ids"].astype(np.int64), |
|
"attention_mask": inputs["attention_mask"].astype(np.int64) |
|
} |
|
|
|
|
|
outputs = session.run(None, onnx_inputs) |
|
|
|
|
|
return { |
|
"embedding": outputs[0].astype(np.float32).tolist(), |
|
"tokens": tokenizer.convert_ids_to_tokens(inputs["input_ids"][0]) |
|
} |
|
|
|
except Exception as e: |
|
raise HTTPException(status_code=400, detail=str(e)) |