feat: expose Cohere and HF embedding support on UI (#236)

Cinnamon · Sep 6, 2024 · 069f0f3 · 069f0f3
1 parent 4d7f164
commit 069f0f3
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 2 deletions.
diff --git a/flowsettings.py b/flowsettings.py
@@ -167,7 +167,7 @@
         "default": False,
     }
 
-    KH_EMBEDDINGS["local-bge-en"] = {
+    KH_EMBEDDINGS["fast_embed"] = {
         "spec": {
             "__type__": "kotaemon.embeddings.FastEmbedEmbeddings",
             "model_name": "BAAI/bge-base-en-v1.5",
@@ -202,6 +202,23 @@
     "default": False,
 }
 
+# additional embeddings configurations
+KH_EMBEDDINGS["cohere"] = {
+    "spec": {
+        "__type__": "kotaemon.embeddings.LCCohereEmbeddings",
+        "model": "embed-multilingual-v2.0",
+        "cohere_api_key": "your-key",
+    },
+    "default": False,
+}
+# KH_EMBEDDINGS["huggingface"] = {
+#     "spec": {
+#         "__type__": "kotaemon.embeddings.LCHuggingFaceEmbeddings",
+#         "model_name": "sentence-transformers/all-mpnet-base-v2",
+#     },
+#     "default": False,
+# }
+
 KH_REASONINGS = [
     "ktem.reasoning.simple.FullQAPipeline",
     "ktem.reasoning.simple.FullDecomposeQAPipeline",

diff --git a/libs/ktem/ktem/embeddings/manager.py b/libs/ktem/ktem/embeddings/manager.py
@@ -55,10 +55,18 @@ def load_vendors(self):
         from kotaemon.embeddings import (
             AzureOpenAIEmbeddings,
             FastEmbedEmbeddings,
+            LCCohereEmbeddings,
+            LCHuggingFaceEmbeddings,
             OpenAIEmbeddings,
         )
 
-        self._vendors = [AzureOpenAIEmbeddings, OpenAIEmbeddings, FastEmbedEmbeddings]
+        self._vendors = [
+            AzureOpenAIEmbeddings,
+            OpenAIEmbeddings,
+            FastEmbedEmbeddings,
+            LCCohereEmbeddings,
+            LCHuggingFaceEmbeddings,
+        ]
 
     def __getitem__(self, key: str) -> BaseEmbeddings:
         """Get model by name"""