summaryrefslogtreecommitdiff
path: root/supplemental/search/fastembed-api/fastembed-server.py
blob: 02da69db26a2130cd8ee13effd0586af157c4781 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
from fastembed import TextEmbedding
from fastapi import FastAPI
from pydantic import BaseModel

models = {}

app = FastAPI()

class EmbeddingRequest(BaseModel):
    model: str
    input: str

@app.post("/v1/embeddings")
def embeddings(request: EmbeddingRequest):
    model = models.get(request.model) or TextEmbedding(request.model)
    models[request.model] = model
    embeddings = next(model.embed(request.input)).tolist()
    return {"data": [{"embedding": embeddings}]}

@app.get("/health")
def health():
    return {"status": "ok"}

if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=11345)