notebooks/llmu/co_aws_ch5_rerank_sm.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a target=\"_blank\" href=\"https://colab.research.google.com/github/cohere-ai/notebooks/blob/main/notebooks/llmu/co_aws_ch5_rerank_sm.ipynb\"> <img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Reranking Using Cohere Rerank on Amazon SageMaker"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Reranking is an essential technique in information retrieval systems, especially in large-scale search applications. It is a process of reordering a set of initially retrieved documents based on their relevance to a user's query.\n",
    "\n",
    "What’s great about reranking is that while it gives a huge boost to search results, implementing Cohere’s Rerank models requires adding only one line of code to any existing search system, whether a semantic search system or a traditional search system that uses keyword-based approaches.\n",
    "\n",
    "In this notebook, we'll explore how to use the Cohere Rerank endpoint on Amazon SageMaker. In particular, we'll look at an example of a multi-aspect search on semi-structured data, and walk through how to perform reranking on email data that contains multiple fields: “title” and “content.”"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Setup"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, let's install and import the necessary libraries and set up our Cohere client.\n",
    "\n",
    "We'll need to create a SageMaker endpoint that exposes access to a Cohere model (Rerank v3 in our case). For this, we’ll use the cohere_aws SDK which makes it easy to set up the endpoint, together with AWS’s boto3 library."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install cohere cohere-aws boto3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import boto3\n",
    "import cohere\n",
    "import cohere_aws\n",
    "from cohere_aws import Client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cohere\n",
    "\n",
    "# Create SageMaker client via the native Cohere SDK\n",
    "# Contact your AWS administrator for the credentials\n",
    "co = cohere.SagemakerClient(\n",
    "    aws_region=\"YOUR_AWS_REGION\",\n",
    "    aws_access_key=\"YOUR_AWS_ACCESS_KEY_ID\",\n",
    "    aws_secret_key=\"YOUR_AWS_SECRET_ACCESS_KEY\",\n",
    "    aws_session_token=\"YOUR_AWS_SESSION_TOKEN\",\n",
    ")\n",
    "\n",
    "# For creating an endpoint, you need to use the cohere_aws client: Set environment variables with the AWS credentials\n",
    "os.environ['AWS_ACCESS_KEY_ID'] = \"YOUR_AWS_ACCESS_KEY_ID\"\n",
    "os.environ['AWS_SECRET_ACCESS_KEY'] = \"YOUR_AWS_SECRET_ACCESS_KEY\"\n",
    "os.environ['AWS_SESSION_TOKEN'] = \"YOUR_AWS_SESSION_TOKEN\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Create Endpoint"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With SageMaker, we’ll need to create an endpoint via an AWS instance. The marketplace listing provides more details, including pricing, on the recommended instance type for a particular model.\n",
    "\n",
    "To create the endpoint, we define:\n",
    "\n",
    "- arn: The model package ARN we defined in the previous step\n",
    "- endpoint_name: A name we can give as an identifier\n",
    "- instance_type: The instance type to be used\n",
    "- n_instances: The number of instances\n",
    "\n",
    "We pass the arguments to the create_endpoint method from the cohere_aws library. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create SageMaker endpoint via the cohere_aws SDK\n",
    "cohere_package = \"cohere-rerank-english-v3-01-d3687e0d2e3a366bb904275616424807\"\n",
    "model_package_map = {\n",
    "    \"us-east-1\": f\"arn:aws:sagemaker:us-east-1:865070037744:model-package/{cohere_package}\",\n",
    "    \"us-east-2\": f\"arn:aws:sagemaker:us-east-2:057799348421:model-package/{cohere_package}\",\n",
    "    \"us-west-1\": f\"arn:aws:sagemaker:us-west-1:382657785993:model-package/{cohere_package}\",\n",
    "    \"us-west-2\": f\"arn:aws:sagemaker:us-west-2:594846645681:model-package/{cohere_package}\",\n",
    "    \"ca-central-1\": f\"arn:aws:sagemaker:ca-central-1:470592106596:model-package/{cohere_package}\",\n",
    "    \"eu-central-1\": f\"arn:aws:sagemaker:eu-central-1:446921602837:model-package/{cohere_package}\",\n",
    "    \"eu-west-1\": f\"arn:aws:sagemaker:eu-west-1:985815980388:model-package/{cohere_package}\",\n",
    "    \"eu-west-2\": f\"arn:aws:sagemaker:eu-west-2:856760150666:model-package/{cohere_package}\",\n",
    "    \"eu-west-3\": f\"arn:aws:sagemaker:eu-west-3:843114510376:model-package/{cohere_package}\",\n",
    "    \"eu-north-1\": f\"arn:aws:sagemaker:eu-north-1:136758871317:model-package/{cohere_package}\",\n",
    "    \"ap-southeast-1\": f\"arn:aws:sagemaker:ap-southeast-1:192199979996:model-package/{cohere_package}\",\n",
    "    \"ap-southeast-2\": f\"arn:aws:sagemaker:ap-southeast-2:666831318237:model-package/{cohere_package}\",\n",
    "    \"ap-northeast-2\": f\"arn:aws:sagemaker:ap-northeast-2:745090734665:model-package/{cohere_package}\",\n",
    "    \"ap-northeast-1\": f\"arn:aws:sagemaker:ap-northeast-1:977537786026:model-package/{cohere_package}\",\n",
    "    \"ap-south-1\": f\"arn:aws:sagemaker:ap-south-1:077584701553:model-package/{cohere_package}\",\n",
    "    \"sa-east-1\": f\"arn:aws:sagemaker:sa-east-1:270155090741:model-package/{cohere_package}\",\n",
    "}\n",
    "\n",
    "region = boto3.Session().region_name\n",
    "\n",
    "if region not in model_package_map.keys():\n",
    "    raise Exception(\"UNSUPPORTED REGION\")\n",
    "\n",
    "model_package_arn = model_package_map[region]\n",
    "\n",
    "co_aws = Client(region_name=region)\n",
    "\n",
    "co_aws.create_endpoint(arn=model_package_arn, endpoint_name=\"my-rerank-v3\", instance_type=\"ml.g5.xlarge\", n_instances=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Retrieve Documents"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let’s assume that the first stage of retrieval has already been performed, whether it’s through a semantic, keyword, or any other type of search system.\n",
    "\n",
    "Here we have a list of nine documents that represent the search results of that first stage. Each document has two fields, Title and Content, corresponding to the contents of an email. Each email is a dictionary containing these fields that preserves its semi-structured format, which the Rerank endpoint can take advantage of."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [],
   "source": [
    "documents = [\n",
    "    {\"Title\":\"Incorrect Password\",\"Content\":\"Hello, I have been trying to access my account for the past hour and it keeps saying my password is incorrect. Can you please help me?\"},\n",
    "    {\"Title\":\"Confirmation Email Missed\",\"Content\":\"Hi, I recently purchased a product from your website but I never received a confirmation email. Can you please look into this for me?\"},\n",
    "    {\"Title\":\"Questions about Return Policy\",\"Content\":\"Hello, I have a question about the return policy for this product. I purchased it a few weeks ago and it is defective.\"},\n",
    "    {\"Title\":\"Customer Support is Busy\",\"Content\":\"Good morning, I have been trying to reach your customer support team for the past week but I keep getting a busy signal. Can you please help me?\"},\n",
    "    {\"Title\":\"Received Wrong Item\",\"Content\":\"Hi, I have a question about my recent order. I received the wrong item and I need to return it.\"},\n",
    "    {\"Title\":\"Customer Service is Unavailable\",\"Content\":\"Hello, I have been trying to reach your customer support team for the past hour but I keep getting a busy signal. Can you please help me?\"},\n",
    "    {\"Title\":\"Return Policy for Defective Product\",\"Content\":\"Hi, I have a question about the return policy for this product. I purchased it a few weeks ago and it is defective.\"},\n",
    "    {\"Title\":\"Wrong Item Received\",\"Content\":\"Good morning, I have a question about my recent order. I received the wrong item and I need to return it.\"},\n",
    "    {\"Title\":\"Return Defective Product\",\"Content\":\"Hello, I have a question about the return policy for this product. I purchased it a few weeks ago and it is defective.\"}\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Rerank Documents"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To use the endpoint, we now use the cohere SDK. Adding a reranking component is simple with Cohere Rerank. It takes just one line of code to implement."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [],
   "source": [
    "query = 'What emails have been about refunds?'\n",
    "\n",
    "response = co.rerank(documents=documents,\n",
    "                     query=query,\n",
    "                     rank_fields=[\"Title\",\"Content\"],\n",
    "                     top_n=3,\n",
    "                     model=\"my-rerank-v3\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# View Results"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since we defined top_n=3, we’ll get the top three most relevant documents to the query. For each document, the response contains the index of its position in the original list and its relevance score against the query."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Documents \n",
      "\n",
      "#1:\n",
      "{'Title': 'Questions about Return Policy', 'Content': 'Hello, I have a question about the return policy for this product. I purchased it a few weeks ago and it is defective.'}\n",
      "\n",
      "#2:\n",
      "{'Title': 'Return Policy for Defective Product', 'Content': 'Hi, I have a question about the return policy for this product. I purchased it a few weeks ago and it is defective.'}\n",
      "\n",
      "#3:\n",
      "{'Title': 'Return Defective Product', 'Content': 'Hello, I have a question about the return policy for this product. I purchased it a few weeks ago and it is defective.'}\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(\"Documents\",\"\\n\")\n",
    "\n",
    "for idx,doc in enumerate(response.results):\n",
    "    print(f\"#{idx+1}:\\n{documents[doc.index]}\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The search query was looking for emails about refunds. But none of the documents mention the word “refunds” specifically.\n",
    "\n",
    "However, the Rerank model was able to retrieve the right documents. From the list of documents, some mention the word “return.” The Rerank model can capture semantically similar meanings between two pieces of text, so it is able to return documents that mention return instead, which has a very similar meaning to return."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Delete endpoint"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Important Note: You will continue to incur charges for as long as an endpoint is running, so remember to delete the endpoint when your usage ends."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "co_aws.delete_endpoint()\n",
    "co_aws.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Reranking is a valuable technique used in information retrieval systems to enhance the relevance of search results. Cohere's Rerank endpoint, including its latest model, Rerank 3, offers improved capabilities for enterprise search.\n",
    "\n",
    "By incorporating reranking with a single line of code, as shown in our example in this chapter, the model successfully identified semantically similar documents, even when specific keywords were absent from the query. This example highlights the potential benefits of integrating reranking into existing search systems to enhance search accuracy and user satisfaction."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "base"
  },
  "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}