Module 3 - Semantic Search.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e87dc259",
   "metadata": {},
   "source": [
    "# Semantic Search with Amazon OpenSearch Service "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b0cfd51d",
   "metadata": {},
   "source": [
    "Now that we've been able to search the data set with a keyword search, let's see how we can use Semantic Search to improve the matches. To do this, we will add a vector representation of the questions to our data set in OpenSearch, then do the same with our sample query \"Does this work with xbox?\". In OpenSearch, we'll use a KNN search to find matches based on a cosine similarity rating on the vector.\n",
    "\n",
    "![word vector](word2vec.png)\n",
    "\n",
    "\n",
    "We will:\n",
    "1. Use a HuggingFace BERT model to generate vector for the PQA dataset\n",
    "2. Upload the dataset to OpenSearch, with the original question and answer text combined with the vector representation of the questions.\n",
    "3. Translate the query question to a vector.\n",
    "4. Perform a KNN search in OpenSearch to perform semantic search"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "31703e3d",
   "metadata": {},
   "source": [
    "### 1. Check PyTorch Version\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ac12126",
   "metadata": {},
   "source": [
    "As in the previous modules, let's import PyTorch and confirm that have have the latest version of PyTorch. The version should already be 2.2.0 or higher. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b532987",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "print(torch.__version__)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f2f1cc51",
   "metadata": {},
   "source": [
    "### 2. Retrieve notebook variables\n",
    "\n",
    "The line below will retrieve your shared variables from the previous notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18a0e06e",
   "metadata": {},
   "outputs": [],
   "source": [
    "%store -r"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0aa614bc",
   "metadata": {},
   "source": [
    "### 3. Import library\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1688f4e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import boto3\n",
    "import re\n",
    "import time\n",
    "import sagemaker"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9510e820",
   "metadata": {},
   "source": [
    "### 4. Prepare BERT Model \n",
    "\n",
    "For this module, we will be using the HuggingFace BERT model to generate vectorization data, where every sentence is 768 dimension data. Let's create some helper functions we'll use later on.\n",
    "![BERT](nlp_bert.png)\n",
    "\n",
    "We are creating 2 functions:\n",
    "1. mean_pooling\n",
    "2. sentence_to_vector - this is the key function we'll use to generate our vector for the headset PQA dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32964815",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from transformers import AutoTokenizer, AutoModel\n",
    "from transformers import DistilBertTokenizer, DistilBertModel\n",
    "\n",
    "#model_name = \"distilbert-base-uncased\"\n",
    "#model_name = \"sentence-transformers/msmarco-distilbert-base-dot-prod-v3\"\n",
    "model_name = \"sentence-transformers/distilbert-base-nli-stsb-mean-tokens\"\n",
    "\n",
    "\n",
    "#Mean Pooling - Take attention mask into account for correct averaging\n",
    "def mean_pooling(model_output, attention_mask):\n",
    "    token_embeddings = model_output[0] #First element of model_output contains all token embeddings\n",
    "    input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()\n",
    "    sum_embeddings = torch.sum(token_embeddings * input_mask_expanded, 1)\n",
    "    sum_mask = torch.clamp(input_mask_expanded.sum(1), min=1e-9)\n",
    "    return sum_embeddings / sum_mask\n",
    "\n",
    "\n",
    "def sentence_to_vector(raw_inputs):\n",
    "    tokenizer = DistilBertTokenizer.from_pretrained(model_name)\n",
    "    model = DistilBertModel.from_pretrained(model_name)\n",
    "    inputs_tokens = tokenizer(raw_inputs, padding=True, return_tensors=\"pt\")\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        outputs = model(**inputs_tokens)\n",
    "\n",
    "    sentence_embeddings = mean_pooling(outputs, inputs_tokens['attention_mask'])\n",
    "    return sentence_embeddings\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c6607721",
   "metadata": {},
   "source": [
    "### 5. Prepare Headset PQA data\n",
    "We have already downloaded the dataset in Module 2, so let's start by ingesting 1000 rows of the data into a Pandas data frame. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b1cf47b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import pandas as pd\n",
    "\n",
    "def load_pqa(file_name,number_rows=1000):\n",
    "    qa_list = []\n",
    "    df = pd.DataFrame(columns=('question', 'answer'))\n",
    "    with open(file_name) as f:\n",
    "        i=0\n",
    "        for line in f:\n",
    "            data = json.loads(line)\n",
    "            df.loc[i] = [data['question_text'],data['answers'][0]['answer_text']]\n",
    "            i+=1\n",
    "            if(i == number_rows):\n",
    "                break\n",
    "    return df\n",
    "\n",
    "\n",
    "qa_list = load_pqa('amazon-pqa/amazon_pqa_headsets.json',number_rows=1000)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "295df3f6",
   "metadata": {},
   "source": [
    "### 6. Convert the text data into vector\n",
    "Using the helper function we created earlier, let's convert the questions from the Headset PQA dataset into vectors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11405a76",
   "metadata": {},
   "outputs": [],
   "source": [
    "vector_sentences = sentence_to_vector(qa_list[\"question\"].tolist())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f54a349",
   "metadata": {},
   "source": [
    "### 7. Create an OpenSearch cluster connection.\n",
    "Next, we'll use Python API to set up connection with OpenSearch Cluster.\n",
    "\n",
    "Note: if you're using a region other than us-east-1, please update the region in the code below. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "405e0e52",
   "metadata": {},
   "outputs": [],
   "source": [
    "from opensearchpy import OpenSearch, RequestsHttpConnection, AWSV4SignerAuth\n",
    "import boto3\n",
    "import json\n",
    "\n",
    "kms = boto3.client('secretsmanager')\n",
    "aos_credentials = json.loads(kms.get_secret_value(SecretId=outputs['OpenSearchSecret'])['SecretString'])\n",
    "\n",
    "region = 'us-east-1' \n",
    "\n",
    "#credentials = boto3.Session().get_credentials()\n",
    "#auth = AWSV4SignerAuth(credentials, region)\n",
    "auth = (aos_credentials['username'], aos_credentials['password'])\n",
    "\n",
    "index_name = 'nlp_pqa'\n",
    "\n",
    "aos_client = OpenSearch(\n",
    "    hosts = [{'host': aos_host, 'port': 443}],\n",
    "    http_auth = auth,\n",
    "    use_ssl = True,\n",
    "    verify_certs = True,\n",
    "    connection_class = RequestsHttpConnection\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "beaabc1e",
   "metadata": {},
   "source": [
    "### 8. Create a index in Amazon Opensearch Service \n",
    "Whereas we previously created an index with 2 fields, this time we'll define the index with 3 fields: the first field ' question_vector' holds the vector representation of the question, the second is the \"question\" for raw sentence and the third field is \"answer\" for the raw answer data.\n",
    "\n",
    "To create the index, we first define the index in JSON, then use the aos_client connection we initiated ealier to create the index in OpenSearch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5eba5754",
   "metadata": {},
   "outputs": [],
   "source": [
    "knn_index = {\n",
    "    \"settings\": {\n",
    "        \"index.knn\": True,\n",
    "        \"index.knn.space_type\": \"cosinesimil\",\n",
    "        \"analysis\": {\n",
    "          \"analyzer\": {\n",
    "            \"default\": {\n",
    "              \"type\": \"standard\",\n",
    "              \"stopwords\": \"_english_\"\n",
    "            }\n",
    "          }\n",
    "        }\n",
    "    },\n",
    "    \"mappings\": {\n",
    "        \"properties\": {\n",
    "            \"question_vector\": {\n",
    "                \"type\": \"knn_vector\",\n",
    "                \"dimension\": 768,\n",
    "                \"store\": True\n",
    "            },\n",
    "            \"question\": {\n",
    "                \"type\": \"text\",\n",
    "                \"store\": True\n",
    "            },\n",
    "            \"answer\": {\n",
    "                \"type\": \"text\",\n",
    "                \"store\": True\n",
    "            }\n",
    "        }\n",
    "    }\n",
    "}\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1330502a",
   "metadata": {},
   "source": [
    "If for any reason you need to recreate your dataset, you can uncomment and execute the following to delete any previously created indexes. If this is the first time you're running this, you can skip this step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a835b9fb",
   "metadata": {},
   "outputs": [],
   "source": [
    "#aos_client.indices.delete(index=\"nlp_pqa\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c6de634d",
   "metadata": {},
   "source": [
    "Using the above index definition, we now need to create the index in Amazon OpenSearch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "715b751d",
   "metadata": {},
   "outputs": [],
   "source": [
    "aos_client.indices.create(index=\"nlp_pqa\",body=knn_index,ignore=400)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a7007735",
   "metadata": {},
   "source": [
    "Let's verify the created index information"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f71659d",
   "metadata": {},
   "outputs": [],
   "source": [
    "aos_client.indices.get(index=\"nlp_pqa\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0040992c",
   "metadata": {},
   "source": [
    "### 9. Load the raw data into the Index\n",
    "Next, let's load the headset enhanced PQA data into the index we've just created."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e55e6a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "i = 0\n",
    "for c in qa_list[\"question\"].tolist():\n",
    "    content=c\n",
    "    vector=vector_sentences[i].tolist()\n",
    "    answer=qa_list[\"answer\"][i]\n",
    "    i+=1\n",
    "    aos_client.index(index='nlp_pqa',body={\"question_vector\": vector, \"question\": content,\"answer\":answer})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "67fad674",
   "metadata": {},
   "source": [
    "To validate the load, we'll query the number of documents number in the index. We should have 1000 hits in the index."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "05ed0b71",
   "metadata": {},
   "outputs": [],
   "source": [
    "res = aos_client.search(index=\"nlp_pqa\", body={\"query\": {\"match_all\": {}}})\n",
    "print(\"Records found: %d.\" % res['hits']['total']['value'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "93888b33",
   "metadata": {},
   "source": [
    "### 10. Generate vector for user input query \n",
    "\n",
    "Next, we'll use the same helper function to translate our input question \"does this work with xbox?\" into a vector. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec91279d",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_raw_sentences = ['does this work with xbox?']\n",
    "search_vector = sentence_to_vector(query_raw_sentences)[0].tolist()\n",
    "search_vector"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "de9b827c",
   "metadata": {},
   "source": [
    "### 11. Search vector with \"Semantic Search\" \n",
    "\n",
    "Now that we have vector in OpenSearch and a vector for our query question, let's perform a KNN search in OpenSearch.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c5f4e81",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "query={\n",
    "    \"size\": 30,\n",
    "    \"query\": {\n",
    "        \"knn\": {\n",
    "            \"question_vector\":{\n",
    "                \"vector\":search_vector,\n",
    "                \"k\":30\n",
    "            }\n",
    "        }\n",
    "    }\n",
    "}\n",
    "\n",
    "res = aos_client.search(index=\"nlp_pqa\", \n",
    "                       body=query,\n",
    "                       stored_fields=[\"question\",\"answer\"])\n",
    "#print(\"Got %d Hits:\" % res['hits']['total']['value'])\n",
    "query_result=[]\n",
    "for hit in res['hits']['hits']:\n",
    "    row=[hit['_id'],hit['_score'],hit['fields']['question'][0],hit['fields']['answer'][0]]\n",
    "    query_result.append(row)\n",
    "\n",
    "query_result_df = pd.DataFrame(data=query_result,columns=[\"_id\",\"_score\",\"question\",\"answer\"])\n",
    "display(query_result_df)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0abddaa4",
   "metadata": {},
   "source": [
    "### 12. Search the same query with \"Text Search\"\n",
    "\n",
    "Let's repeat the same query with a keyword search and compare the differences."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c652c52",
   "metadata": {},
   "outputs": [],
   "source": [
    "query={\n",
    "    \"size\": 30,\n",
    "    \"query\": {\n",
    "        \"match\": {\n",
    "            \"question\":\"does this work with xbox?\"\n",
    "        }\n",
    "    }\n",
    "}\n",
    "\n",
    "res = aos_client.search(index=\"nlp_pqa\", \n",
    "                       body=query,\n",
    "                       stored_fields=[\"question\",\"answer\"])\n",
    "#print(\"Got %d Hits:\" % res['hits']['total']['value'])\n",
    "query_result=[]\n",
    "for hit in res['hits']['hits']:\n",
    "    row=[hit['_id'],hit['_score'],hit['fields']['question'][0],hit['fields']['answer'][0]]\n",
    "    query_result.append(row)\n",
    "\n",
    "query_result_df = pd.DataFrame(data=query_result,columns=[\"_id\",\"_score\",\"question\",\"answer\"])\n",
    "display(query_result_df)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb777d3d",
   "metadata": {},
   "source": [
    "### 13. Observe The Results\n",
    "\n",
    "Compare the first few records in the two searches above. For the Semantic search, the first 10 or so results are very similar to our input questions, as we expect. Compare this to keyword search, where the results quickly start to deviate from our search query (e.g. \"it shows xbox 360. Does it work for ps3 as well?\" - this matches on keywords but has a different meaning)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "607e909b",
   "metadata": {},
   "source": [
    "### 14. Store Variables Used for the Next Notebook\n",
    "\n",
    "There are a few values you will need for the next notebook, execute the cells below to store them so they can be copied and pasted into the next part of the exercise."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8db65388",
   "metadata": {},
   "outputs": [],
   "source": [
    "%store qa_list"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda_pytorch_p310",
   "language": "python",
   "name": "conda_pytorch_p310"
  },
  "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.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}