summaryrefslogtreecommitdiff
path: root/supplemental/search/fastembed-api/fastembed-server.py
diff options
context:
space:
mode:
authorLain Soykaf <lain@lain.com>2024-05-28 16:51:19 +0400
committerLain Soykaf <lain@lain.com>2024-05-28 16:51:19 +0400
commitcc42b50c5bbbeee066d3d26907eebb488ec6aa3d (patch)
tree3b67771624d237bdb58f0931098d6682191cc214 /supplemental/search/fastembed-api/fastembed-server.py
parentfdc3cbb8cbefad2161cc408b4440420d6e4b2a88 (diff)
parentbef15cde6141a977aebdfc998d6091a31c4fc2d6 (diff)
downloadpleroma-cc42b50c5bbbeee066d3d26907eebb488ec6aa3d.tar.gz
pleroma-cc42b50c5bbbeee066d3d26907eebb488ec6aa3d.zip
Merge branch 'develop' of git.pleroma.social:pleroma/pleroma into pleroma-feature/akkoma-prune-old-posts
Diffstat (limited to 'supplemental/search/fastembed-api/fastembed-server.py')
-rw-r--r--supplemental/search/fastembed-api/fastembed-server.py27
1 files changed, 27 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..02da69db2
--- /dev/null
+++ b/supplemental/search/fastembed-api/fastembed-server.py
@@ -0,0 +1,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)