pleroma/supplemental/search/fastembed-api/fastembed-server.py

28 lines
635 B
Python
Raw Permalink Normal View History

2024-05-18 00:55:17 -07:00
from fastembed import TextEmbedding
from fastapi import FastAPI
from pydantic import BaseModel
models = {}
2024-05-18 00:55:17 -07:00
app = FastAPI()
class EmbeddingRequest(BaseModel):
model: str
input: str
2024-05-18 00:55:17 -07:00
@app.post("/v1/embeddings")
2024-05-18 00:55:17 -07:00
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}]}
2024-05-18 00:55:17 -07:00
@app.get("/health")
def health():
return {"status": "ok"}
2024-05-18 00:55:17 -07:00
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=11345)