embeddings / app.py
AWeirdDev's picture
Update app.py
5e25768 verified
raw
history blame contribute delete
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)
)