Merge branch 'main' into embeddings

Author: trducng

libs/kotaemon/kotaemon/embeddings/__init__.py

@@ -1,5 +1,6 @@
 from .base import BaseEmbeddings
 from .endpoint_based import EndpointEmbeddings
+from .fastembed import FastEmbedEmbeddings
 from .langchain_based import (
     LCAzureOpenAIEmbeddings,
     LCCohereEmbdeddings,
@@ -17,4 +18,5 @@
     "LCHuggingFaceEmbeddings",
     "OpenAIEmbeddings",
     "AzureOpenAIEmbeddings",
+    "FastEmbedEmbeddings",
 ]

libs/kotaemon/kotaemon/embeddings/base.py

@@ -18,3 +18,12 @@ async def ainvoke(
         self, text: str | list[str] | Document | list[Document], *args, **kwargs
     ) -> list[DocumentWithEmbedding]:
         raise NotImplementedError
+
+    def prepare_input(
+        self, text: str | list[str] | Document | list[Document]
+    ) -> list[Document]:
+        if isinstance(text, (str, Document)):
+            return [Document(content=text)]
+        elif isinstance(text, list):
+            return [Document(content=_) for _ in text]
+        return text

libs/kotaemon/kotaemon/embeddings/fastembed.py

@@ -0,0 +1,68 @@
+from typing import TYPE_CHECKING, Optional
+
+from kotaemon.base import Document, DocumentWithEmbedding, Param
+
+from .base import BaseEmbeddings
+
+if TYPE_CHECKING:
+    from fastembed import TextEmbedding
+
+
+class FastEmbedEmbeddings(BaseEmbeddings):
+    """Utilize fastembed library for embeddings locally without GPU.
+
+    Supported model: https://qdrant.github.io/fastembed/examples/Supported_Models/
+    Code: https://github.com/qdrant/fastembed
+    """
+
+    model_name: str = Param(
+        "BAAI/bge-small-en-v1.5",
+        help=(
+            "Model name for fastembed. "
+            "Supported model: "
+            "https://qdrant.github.io/fastembed/examples/Supported_Models/"
+        ),
+    )
+    batch_size: int = Param(
+        256,
+        help="Batch size for embeddings. Higher values use more memory, but are faster",
+    )
+    parallel: Optional[int] = Param(
+        None,
+        help=(
+            "Number of threads to use for embeddings. "
+            "If > 1, data-parallel encoding will be used. "
+            "If 0, use all available CPUs. "
+            "If None, use default onnxruntime threading. "
+            "Defaults to None"
+        ),
+    )
+
+    @Param.auto()
+    def client_(self) -> "TextEmbedding":
+        from fastembed import TextEmbedding
+
+        return TextEmbedding(model_name=self.model_name)
+
+    def invoke(
+        self, text: str | list[str] | Document | list[Document], *args, **kwargs
+    ) -> list[DocumentWithEmbedding]:
+        input_ = self.prepare_input(text)
+        embeddings = self.client_.embed(
+            [_.content for _ in input_],
+            batch_size=self.batch_size,
+            parallel=self.parallel,
+        )
+        return [
+            DocumentWithEmbedding(
+                content=doc,
+                embedding=list(embedding),
+            )
+            for doc, embedding in zip(input_, embeddings)
+        ]
+
+    async def ainvoke(
+        self, text: str | list[str] | Document | list[Document], *args, **kwargs
+    ) -> list[DocumentWithEmbedding]:
+        """Fastembed does not support async API."""
+        return self.invoke(text, *args, **kwargs)

libs/kotaemon/kotaemon/embeddings/openai.py

@@ -41,15 +41,6 @@ def max_retries_(self):
             return DEFAULT_MAX_RETRIES
         return self.max_retries
 
-    def prepare_input(
-        self, text: str | list[str] | Document | list[Document]
-    ) -> list[Document]:
-        if isinstance(text, (str, Document)):
-            return [Document(content=text)]
-        elif isinstance(text, list):
-            return [Document(content=_) for _ in text]
-        return text
-
     def prepare_client(self, async_version: bool = False):
         """Get the OpenAI client
 

libs/kotaemon/pyproject.toml

@@ -11,7 +11,7 @@ packages.find.exclude = ["tests*", "env*"]
 # metadata and dependencies
 [project]
 name = "kotaemon"
-version = "0.3.9"
+version = "0.3.10"
 requires-python = ">= 3.10"
 description = "Kotaemon core library for AI development."
 dependencies = [
@@ -61,6 +61,7 @@ adv = [
     "elasticsearch",
     "llama-cpp-python",
     "pdfservices-sdk @  git+https://github.com/niallcm/pdfservices-python-sdk.git@bump-and-unfreeze-requirements",
+    "fastembed",
 ]
 dev = [
     "ipython",

libs/kotaemon/tests/test_embedding_models.py

@@ -7,6 +7,7 @@
 from kotaemon.base import Document
 from kotaemon.embeddings import (
     AzureOpenAIEmbeddings,
+    FastEmbedEmbeddings,
     LCAzureOpenAIEmbeddings,
     LCCohereEmbdeddings,
     LCHuggingFaceEmbeddings,
@@ -20,6 +21,13 @@
     openai_embedding = CreateEmbeddingResponse.model_validate(json.load(f))
 
 
+def assert_embedding_result(output):
+    assert isinstance(output, list)
+    assert isinstance(output[0], Document)
+    assert isinstance(output[0].embedding, list)
+    assert isinstance(output[0].embedding[0], float)
+
+
 @patch(
     "openai.resources.embeddings.Embeddings.create",
     side_effect=lambda *args, **kwargs: openai_embedding,
@@ -32,10 +40,7 @@ def test_lcazureopenai_embeddings_raw(openai_embedding_call):
         openai_api_key="some-key",
     )
     output = model("Hello world")
-    assert isinstance(output, list)
-    assert isinstance(output[0], Document)
-    assert isinstance(output[0].embedding, list)
-    assert isinstance(output[0].embedding[0], float)
+    assert_embedding_result(output)
     openai_embedding_call.assert_called()
 
 
@@ -51,10 +56,67 @@ def test_lcazureopenai_embeddings_batch_raw(openai_embedding_call):
         openai_api_key="some-key",
     )
     output = model(["Hello world", "Goodbye world"])
-    assert isinstance(output, list)
-    assert isinstance(output[0], Document)
-    assert isinstance(output[0].embedding, list)
-    assert isinstance(output[0].embedding[0], float)
+    assert_embedding_result(output)
+    openai_embedding_call.assert_called()
+
+
+@patch(
+    "openai.resources.embeddings.Embeddings.create",
+    side_effect=lambda *args, **kwargs: openai_embedding,
+)
+def test_azureopenai_embeddings_raw(openai_embedding_call):
+    model = AzureOpenAIEmbeddings(
+        azure_endpoint="https://test.openai.azure.com/",
+        api_key="some-key",
+        api_version="version",
+        azure_deployment="text-embedding-ada-002",
+    )
+    output = model("Hello world")
+    assert_embedding_result(output)
+    openai_embedding_call.assert_called()
+
+
+@patch(
+    "openai.resources.embeddings.Embeddings.create",
+    side_effect=lambda *args, **kwargs: openai_embedding_batch,
+)
+def test_azureopenai_embeddings_batch_raw(openai_embedding_call):
+    model = AzureOpenAIEmbeddings(
+        azure_endpoint="https://test.openai.azure.com/",
+        api_key="some-key",
+        api_version="version",
+        azure_deployment="text-embedding-ada-002",
+    )
+    output = model(["Hello world", "Goodbye world"])
+    assert_embedding_result(output)
+    openai_embedding_call.assert_called()
+
+
+@patch(
+    "openai.resources.embeddings.Embeddings.create",
+    side_effect=lambda *args, **kwargs: openai_embedding,
+)
+def test_openai_embeddings_raw(openai_embedding_call):
+    model = OpenAIEmbeddings(
+        api_key="some-key",
+        model="text-embedding-ada-002",
+    )
+    output = model("Hello world")
+    assert_embedding_result(output)
+    openai_embedding_call.assert_called()
+
+
+@patch(
+    "openai.resources.embeddings.Embeddings.create",
+    side_effect=lambda *args, **kwargs: openai_embedding_batch,
+)
+def test_openai_embeddings_batch_raw(openai_embedding_call):
+    model = OpenAIEmbeddings(
+        api_key="some-key",
+        model="text-embedding-ada-002",
+    )
+    output = model(["Hello world", "Goodbye world"])
+    assert_embedding_result(output)
     openai_embedding_call.assert_called()
 
 
@@ -148,10 +210,7 @@ def test_lchuggingface_embeddings(
     )
 
     output = model("Hello World")
-    assert isinstance(output, list)
-    assert isinstance(output[0], Document)
-    assert isinstance(output[0].embedding, list)
-    assert isinstance(output[0].embedding[0], float)
+    assert_embedding_result(output)
     sentence_transformers_init.assert_called()
     langchain_huggingface_embedding_call.assert_called()
 
@@ -166,8 +225,11 @@ def test_lccohere_embeddings(langchain_cohere_embedding_call):
     )
 
     output = model("Hello World")
-    assert isinstance(output, list)
-    assert isinstance(output[0], Document)
-    assert isinstance(output[0].embedding, list)
-    assert isinstance(output[0].embedding[0], float)
+    assert_embedding_result(output)
     langchain_cohere_embedding_call.assert_called()
+
+
+def test_fastembed_embeddings():
+    model = FastEmbedEmbeddings()
+    output = model("Hello World")
+    assert_embedding_result(output)

libs/ktem/flowsettings.py

@@ -57,7 +57,7 @@
     if config("AZURE_OPENAI_EMBEDDINGS_DEPLOYMENT", default=""):
         KH_EMBEDDINGS["azure"] = {
             "spec": {
-                "__type__": "kotaemon.embeddings.AzureOpenAIEmbeddings",
+                "__type__": "kotaemon.embeddings.LCAzureOpenAIEmbeddings",
                 "azure_endpoint": config("AZURE_OPENAI_ENDPOINT", default=""),
                 "api_key": config("AZURE_OPENAI_API_KEY", default=""),
                 "api_version": config("OPENAI_API_VERSION", default="")
@@ -88,7 +88,7 @@
     if len(KH_EMBEDDINGS) < 1:
         KH_EMBEDDINGS["openai"] = {
             "spec": {
-                "__type__": "kotaemon.embeddings.OpenAIEmbeddings",
+                "__type__": "kotaemon.embeddings.LCOpenAIEmbeddings",
                 "base_url": config("OPENAI_API_BASE", default="")
                 or "https://api.openai.com/v1",
                 "api_key": config("OPENAI_API_KEY", default=""),