Merge pull request #90 from premAI-io/hot-fix/embeddings

added support for string in embeddings
premAI-io · Aug 24, 2023 · 719bf08 · 719bf08
2 parents 8e4e2a3 + f78bff9
commit 719bf08
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/ebd-all-minilm/build.sh b/ebd-all-minilm/build.sh
@@ -1,6 +1,6 @@
 #!/bin/bash
 set -e
-export VERSION=1.0.2
+export VERSION=1.0.3
 source "$(dirname "${BASH_SOURCE[0]}")/../utils.sh"
 
 build_cpu ghcr.io/premai-io/embeddings-all-minilm-l6-v2-cpu all-MiniLM-L6-v2 ${@:1}

diff --git a/ebd-all-minilm/routes.py b/ebd-all-minilm/routes.py
@@ -8,7 +8,7 @@
 
 class EmbeddingsInput(BaseModel):
     model: str = None
-    input: Union[List[str], List[List[int]]]
+    input: Union[List[str], List[List[int]], str]
     user: str = ""
 
 
@@ -44,7 +44,10 @@ async def health():
 
 @router.post("/embeddings", response_model=EmbeddingsResponse)
 async def embeddings(body: EmbeddingsInput):
-    values = model.embeddings(texts=body.input)
+    if isinstance(body.input, str):
+        values = model.embeddings(texts=[body.input])
+    else:
+        values = model.embeddings(texts=body.input)
     return EmbeddingsResponse(
         object="list",
         data=[EmbeddingObject(embedding=value) for value in values],