embeddings / app.py
AWeirdDev's picture
Update app.py
5e25768 verified
raw
history blame
517 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():
return embedding_model.encode(req.input).tolist()
return JSONResponse(
await asyncio.to_thread(do)
)