Add Cohere vectorizer (#98)

Adding support for CohereTextVectorizer with embed and embed_many
functionality. Contributed by @nabsabraham

---------

Co-authored-by: nabsabraham <nabila.abraham@gmail.com>

Author: 2 people

.github/workflows/run_tests.yml

@@ -59,6 +59,7 @@ jobs:
         OPENAI_API_KEY: ${{ secrets.OPENAI_KEY }}
         GCP_LOCATION: ${{ secrets.GCP_LOCATION }}
         GCP_PROJECT_ID: ${{ secrets.GCP_PROJECT_ID }}
+        COHERE_API_KEY: ${{ secrets.COHERE_API_KEY }}
       run: |
         make test-cov
 
@@ -73,6 +74,7 @@ jobs:
         OPENAI_API_KEY: ${{ secrets.OPENAI_KEY }}
         GCP_LOCATION: ${{ secrets.GCP_LOCATION }}
         GCP_PROJECT_ID: ${{ secrets.GCP_PROJECT_ID }}
+        COHERE_API_KEY: ${{ secrets.COHERE_API_KEY }}
       run: |
         cd docs/ && treon -v --exclude="./examples/openai_qna.ipynb"
 

conftest.py

@@ -26,6 +26,10 @@ def client():
 def openai_key():
     return os.getenv("OPENAI_API_KEY")
 
+@pytest.fixture
+def cohere_key():
+    return os.getenv("COHERE_API_KEY")
+
 @pytest.fixture
 def gcp_location():
     return os.getenv("GCP_LOCATION")

docs/api/vectorizer.rst

@@ -61,3 +61,22 @@ VertexAITextVectorizer
    :inherited-members:
    :members:
 
+
+CohereTextVectorizer
+======================
+
+.. _coheretextvectorizer_api:
+
+.. currentmodule:: redisvl.vectorize.text.cohere
+
+.. autosummary::
+
+    CohereTextVectorizer.__init__
+    CohereTextVectorizer.embed
+    CohereTextVectorizer.embed_many
+
+.. autoclass:: CohereTextVectorizer
+   :show-inheritance:
+   :inherited-members:
+   :members:
+

docs/user_guide/vectorizers_04.ipynb

@@ -11,6 +11,7 @@
     "1. OpenAI\n",
     "2. HuggingFace\n",
     "3. Vertex AI\n",
+    "4. Cohere\n",
     "\n",
     "Before running this notebook, be sure to\n",
     "1. Have installed ``redisvl`` and have that environment active for this notebook.\n",
@@ -27,7 +28,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": 2,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -298,6 +299,76 @@
     "test[:10]"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Cohere\n",
+    "\n",
+    "[Cohere](https://dashboard.cohere.ai/) allows you to implement language AI into your product. The `CohereTextVectorizer` makes it simple to use RedisVL with the embeddings models at Cohere. For this you will need to install `cohere`.\n",
+    "\n",
+    "```bash\n",
+    "pip install cohere\n",
+    "```"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import getpass\n",
+    "# setup the API Key\n",
+    "api_key = os.environ.get(\"COHERE_API_KEY\") or getpass.getpass(\"Enter your Cohere API key: \")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "\n",
+    "Special attention needs to be paid to the `input_type` parameter for each `embed` call. For example, for embedding \n",
+    "queries, you should set `input_type='search_query'`; for embedding documents, set `input_type='search_document'`. See\n",
+    "more information [here](https://docs.cohere.com/reference/embed)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Vector dimensions:  1024\n",
+      "[-0.010856628, -0.019683838, -0.0062179565, 0.003545761, -0.047943115, 0.0009365082, -0.005924225, 0.016174316, -0.03289795, 0.049194336]\n",
+      "Vector dimensions:  1024\n",
+      "[-0.010108948, -0.016693115, -0.0002310276, -0.022644043, -0.04147339, 0.0021324158, -0.033477783, -0.0005378723, -0.02619934, 0.058013916]\n"
+     ]
+    }
+   ],
+   "source": [
+    "from redisvl.vectorize.text import CohereTextVectorizer\n",
+    "\n",
+    "# create a vectorizer\n",
+    "co = CohereTextVectorizer(\n",
+    "    model=\"embed-english-v3.0\",\n",
+    "    api_config={\"api_key\": api_key},\n",
+    ")\n",
+    "\n",
+    "# embed a search query\n",
+    "test = co.embed(\"This is a test sentence.\", input_type='search_query')\n",
+    "print(\"Vector dimensions: \", len(test))\n",
+    "print(test[:10])\n",
+    "\n",
+    "# embed a document\n",
+    "test = co.embed(\"This is a test sentence.\", input_type='search_document')\n",
+    "print(\"Vector dimensions: \", len(test))\n",
+    "print(test[:10])"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},

redisvl/vectorize/base.py

@@ -28,6 +28,7 @@ def embed_many(
         preprocess: Optional[Callable] = None,
         batch_size: int = 1000,
         as_buffer: bool = False,
+        **kwargs,
     ) -> List[List[float]]:
         raise NotImplementedError
 
@@ -36,6 +37,7 @@ def embed(
         text: str,
         preprocess: Optional[Callable] = None,
         as_buffer: bool = False,
+        **kwargs,
     ) -> List[float]:
         raise NotImplementedError
 
@@ -45,6 +47,7 @@ async def aembed_many(
         preprocess: Optional[Callable] = None,
         batch_size: int = 1000,
         as_buffer: bool = False,
+        **kwargs,
     ) -> List[List[float]]:
         raise NotImplementedError
 
@@ -53,6 +56,7 @@ async def aembed(
         text: str,
         preprocess: Optional[Callable] = None,
         as_buffer: bool = False,
+        **kwargs,
     ) -> List[float]:
         raise NotImplementedError
 

redisvl/vectorize/text/__init__.py

@@ -1,5 +1,11 @@
+from redisvl.vectorize.text.cohere import CohereTextVectorizer
 from redisvl.vectorize.text.huggingface import HFTextVectorizer
 from redisvl.vectorize.text.openai import OpenAITextVectorizer
 from redisvl.vectorize.text.vertexai import VertexAITextVectorizer
 
-__all__ = ["OpenAITextVectorizer", "HFTextVectorizer", "VertexAITextVectorizer"]
+__all__ = [
+    "OpenAITextVectorizer",
+    "HFTextVectorizer",
+    "VertexAITextVectorizer",
+    "CohereTextVectorizer",
+]