import asyncio
from fastapi import FastAPI
from fastapi.responses import JSONResponse
from pydantic import BaseModel

from transformers import AutoModel

embedding_model = AutoModel.from_pretrained('jinaai/jina-embeddings-v2-base-en', trust_remote_code=True)

app = FastAPI()

class Req(BaseModel):
    input: list[str]

@app.post("/embeddings")
async def embeddings(req: Req):
    def do():
        return embedding_model.encode(req.input).tolist()
    return JSONResponse(
        await asyncio.to_thread(do)
    )