demo-python/code/community-integration/cohere/azure-search-cohere-embed-v3-sample.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Azure AI Search with Cohere Embed V3 Int8 Support\n",
    "This code demonstrates the use of the Cohere API to generate embeddings with the latest and highest-performing model from Cohere, [Cohere Embed V3](https://txt.cohere.com/introducing-embed-v3/). It also explains how to store these embeddings in Azure AI Search using the Python SDK as an Int8 vector data type. This approach maintains 100% of the embedding quality while reducing our index size by 4x compared to the Float32 representation.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set up a Python virtual environment in Visual Studio Code\n",
    "\n",
    "1. Open the Command Palette (Ctrl+Shift+P).\n",
    "1. Search for **Python: Create Environment**.\n",
    "1. Select **Venv**.\n",
    "1. Select a Python interpreter. Choose 3.10 or later.\n",
    "\n",
    "It can take a minute to set up. If you run into problems, see [Python environments in VS Code](https://code.visualstudio.com/docs/python/environments)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Install required libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install -r azure-search-cohere-embed-v3-sample-requirements.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cohere\n",
    "import numpy as np\n",
    "import os\n",
    "from dotenv import load_dotenv\n",
    "from azure.search.documents.indexes import SearchIndexClient\n",
    "from azure.search.documents import SearchClient\n",
    "from azure.search.documents.models import (\n",
    "    VectorizedQuery,\n",
    ")\n",
    "from azure.search.documents.indexes.models import (\n",
    "    HnswAlgorithmConfiguration,\n",
    "    SearchField,\n",
    "    SearchableField,\n",
    "    SearchFieldDataType,\n",
    "    SearchIndex,\n",
    "    SimpleField,\n",
    "    VectorSearch,\n",
    "    VectorSearchAlgorithmKind,\n",
    "    VectorSearchProfile,\n",
    ")\n",
    "from azure.core.credentials import AzureKeyCredential\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set Up Cohere and Azure Credentials\n",
    "Before generating embeddings or interacting with Azure AI Search, we need to set up our credentials for both Cohere and Azure AI Search."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "load_dotenv()\n",
    "cohere_api_key = os.getenv(\"COHERE_API_KEY\")\n",
    "co = cohere.Client(cohere_api_key)\n",
    "\n",
    "search_service_endpoint = os.getenv(\"AZURE_SEARCH_SERVICE_ENDPOINT\")\n",
    "search_service_api_key = os.getenv(\"AZURE_SEARCH_ADMIN_KEY\")\n",
    "index_name = \"cohere-embed-v3-index\"\n",
    "credential = AzureKeyCredential(search_service_api_key)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generate Embeddings Function\n",
    "This function will use the Cohere API to generate int8 embeddings for a list of documents. These embeddings are optimized for search document use cases."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_embeddings(texts, input_type=\"search_document\"):\n",
    "    model = \"embed-english-v3.0\"\n",
    "    # Ensure texts is a list\n",
    "    if isinstance(texts, str):\n",
    "        texts = [texts]\n",
    "\n",
    "    response = co.embed(\n",
    "        texts=texts,\n",
    "        model=model,\n",
    "        input_type=input_type,\n",
    "        embedding_types=[\"int8\"],\n",
    "    )\n",
    "    return [embedding for embedding in response.embeddings.int8]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create or Update Azure AI Search Index\n",
    "This function creates or updates an Azure AI Search index to include a vector field for storing the document embeddings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_or_update_index(client, index_name):\n",
    "    fields = [\n",
    "        SimpleField(name=\"id\", type=SearchFieldDataType.String, key=True),\n",
    "        SearchField(\n",
    "            name=\"text\",\n",
    "            type=SearchFieldDataType.String,\n",
    "            searchable=True,\n",
    "        ),\n",
    "        SearchField(\n",
    "            name=\"embedding\",\n",
    "            type=\"Collection(Edm.SByte)\",  # OData syntax for 8-bit signed integer\n",
    "            vector_search_dimensions=1024,\n",
    "            vector_search_profile_name=\"my-vector-config\",\n",
    "            # hidden=False, Use hidden=False if you want to return the embeddings in the search results\n",
    "        ),\n",
    "    ]\n",
    "\n",
    "    vector_search = VectorSearch(\n",
    "        profiles=[\n",
    "            VectorSearchProfile(\n",
    "                name=\"my-vector-config\",\n",
    "                algorithm_configuration_name=\"my-hnsw\",\n",
    "            )\n",
    "        ],\n",
    "        algorithms=[\n",
    "            HnswAlgorithmConfiguration(\n",
    "                name=\"my-hnsw\",\n",
    "                kind=VectorSearchAlgorithmKind.HNSW,\n",
    "            )\n",
    "        ],\n",
    "    )\n",
    "\n",
    "    index = SearchIndex(name=index_name, fields=fields, vector_search=vector_search)\n",
    "    client.create_or_update_index(index=index)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Index Documents and Their Embeddings\n",
    "Finally, this function indexes the documents along with their int8 embeddings into Azure AI Search. For demonstration, document IDs are generated sequentially, but in practical applications, it's essential to use meaningful identifiers like database row ID, unique filenames, or any other unique metadata associated with the document.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "def index_documents(search_client, documents, embeddings):\n",
    "    documents_to_index = [\n",
    "        {\"id\": str(idx), \"text\": doc, \"embedding\": emb}\n",
    "        for idx, (doc, emb) in enumerate(zip(documents, embeddings))\n",
    "    ]\n",
    "    search_client.upload_documents(documents=documents_to_index)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run the workflow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "documents = [\n",
    "    \"Alan Turing  was an English mathematician, computer scientist, logician, cryptanalyst, philosopher and theoretical biologist.\",\n",
    "    \"Albert Einstein was a German-born theoretical physicist who is widely held to be one of the greatest and most influential scientists of all time.\",\n",
    "    \"Isaac Newton was an English polymath active as a mathematician, physicist, astronomer, alchemist, theologian, and author who was described in his time as a natural philosopher.\",\n",
    "    \"Marie Curie was a Polish and naturalised-French physicist and chemist who conducted pioneering research on radioactivity\"\n",
    "]\n",
    "\n",
    "# Generate embeddings\n",
    "embeddings = generate_embeddings(documents)\n",
    "\n",
    "# Initialize Azure Search Index Client\n",
    "search_index_client = SearchIndexClient(\n",
    "    endpoint=search_service_endpoint,\n",
    "    credential=credential,\n",
    "    index_name=index_name\n",
    ")\n",
    "\n",
    "# Create or update the search index to include the embedding field\n",
    "create_or_update_index(search_index_client, index_name)\n",
    "\n",
    "# Initialize the SearchClient\n",
    "search_client = SearchClient(\n",
    "    endpoint=search_service_endpoint, \n",
    "    index_name=index_name, \n",
    "    credential=credential\n",
    ")\n",
    "\n",
    "# Index the documents and their embeddings\n",
    "index_documents(search_client, documents, embeddings)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Perform a Vector Search"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Title: Alan Turing  was an English mathematician, computer scientist, logician, cryptanalyst, philosopher and theoretical biologist.\n",
      "Score: 0.6248218\n",
      "\n",
      "Title: Albert Einstein was a German-born theoretical physicist who is widely held to be one of the greatest and most influential scientists of all time.\n",
      "Score: 0.5913683\n",
      "\n",
      "Title: Marie Curie was a Polish and naturalised-French physicist and chemist who conducted pioneering research on radioactivity\n",
      "Score: 0.57643753\n",
      "\n"
     ]
    }
   ],
   "source": [
    "from azure.search.documents import SearchClient\n",
    "\n",
    "# Query for vector search\n",
    "query = \"foundational figures in computer science\"\n",
    "\n",
    "# Generate query embeddings\n",
    "# Use input_type=\"search_query\" for query embeddings\n",
    "query_embeddings = generate_embeddings(query, input_type=\"search_query\")\n",
    "\n",
    "search_client = SearchClient(search_service_endpoint, index_name, credential)\n",
    "\n",
    "vector_query = VectorizedQuery(\n",
    "    vector=query_embeddings[0], k_nearest_neighbors=3, fields=\"embedding\"\n",
    ")\n",
    "\n",
    "results = search_client.search(\n",
    "    search_text=None,  # No search text for pure vector search\n",
    "    vector_queries=[vector_query],\n",
    ")\n",
    "\n",
    "for result in results:\n",
    "    print(f\"Text: {result['text']}\")\n",
    "    print(f\"Score: {result['@search.score']}\\n\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}