diff options
Diffstat (limited to 'supplemental/search/fastembed-api/fastembed-server.py')
-rw-r--r-- | supplemental/search/fastembed-api/fastembed-server.py | 23 |
1 files changed, 23 insertions, 0 deletions
diff --git a/supplemental/search/fastembed-api/fastembed-server.py b/supplemental/search/fastembed-api/fastembed-server.py new file mode 100644 index 000000000..dd4a7a9c8 --- /dev/null +++ b/supplemental/search/fastembed-api/fastembed-server.py @@ -0,0 +1,23 @@ +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}]} + +if __name__ == "__main__": + import uvicorn + + uvicorn.run(app, host="0.0.0.0", port=11345) |