pleroma

commit: b9af017a4cf1025c7d8245fa4f1dbcb678ddd4b9
parent 72ec261a69a7dda7ab95667e425824ab7758b636
Author: Lain Soykaf <lain@lain.com>
Date:   Sun, 19 May 2024 12:33:49 +0400

B FastembedServer: Switch to OpenAI api, support changing models

Diffstat:
M
python/fastembed-server.py
12
+++++++-----

1 file changed, 7 insertions(+), 5 deletions(-)
diff --git a/python/fastembed-server.py b/python/fastembed-server.py
@@ -2,18 +2,20 @@ from fastembed import TextEmbedding
 from fastapi import FastAPI
 from pydantic import BaseModel
 
-model = TextEmbedding("snowflake/snowflake-arctic-embed-xs")
+models = {}
 
 app = FastAPI()
 
 class EmbeddingRequest(BaseModel):
     model: str
-    prompt: str
+    input: str
 
-@app.post("/api/embeddings")
+@app.post("/v1/embeddings")
 def embeddings(request: EmbeddingRequest):
-    embeddings = next(model.embed(request.prompt)).tolist()
-    return {"embedding": embeddings}
+    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