embeddings / app.py
AWeirdDev's picture
Create app.py
0d60a24 verified
raw
history blame
510 Bytes
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():
embedding_model.encode(req.input).tolist()
return JSONResponse(
await asyncio.to_thread(do)
)