Cohere embeddings support (#432)

cohere embedding support Co-authored-by: Abhinav Naikawadi <[email protected]>
refuel-ai · Jul 6, 2023 · d568318 · d568318
1 parent 4250884
commit d568318
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/src/autolabel/few_shot/__init__.py b/src/autolabel/few_shot/__init__.py
@@ -4,6 +4,7 @@
 from autolabel.configs import AutolabelConfig
 from autolabel.schema import FewShotAlgorithm, ModelProvider
 from langchain.embeddings import (
+    CohereEmbeddings,
     HuggingFaceEmbeddings,
     OpenAIEmbeddings,
     VertexAIEmbeddings,
@@ -30,6 +31,7 @@
     ModelProvider.OPENAI: OpenAIEmbeddings,
     ModelProvider.GOOGLE: VertexAIEmbeddings,
     ModelProvider.HUGGINGFACE_PIPELINE: HuggingFaceEmbeddings,
+    ModelProvider.COHERE: CohereEmbeddings,
 }
 
 logger = logging.getLogger(__name__)