Module 5 - Semantic Search with Fine Tuned Model.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "256115b6",
   "metadata": {},
   "source": [
    "# Semantic Search with Fine Tuning Model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "47f3e6d7",
   "metadata": {},
   "source": [
    "In the Module 3 and Module 4, we are using pretrained BERT model to convert text into vector. However, to improve the model accuary, we need to fine tune the BERT model. See [Fine-tune a pretrained model](https://huggingface.co/docs/transformers/training) for more informaiton.\n",
    "\n",
    "In this module, we will use [Amazon Product Question and Answer (PQA) dataset](https://registry.opendata.aws/amazon-pqa/) to fine tune the model. The architecture of this module is like follows:\n",
    "\n",
    "![semantic_search_with_fine_tuning](semantic_search_with_fine_tuning.png)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "533b103a",
   "metadata": {},
   "source": [
    "### 1.Import PyTorch and check version.\n",
    "\n",
    "As in the previous modules, let's import PyTorch and confirm that have have the latest version of PyTorch. The version should already be 1.10.2 or higher. If not, please run the lab in order to get everything set up."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d0b2ace6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "print(torch.__version__)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b594ab8d",
   "metadata": {},
   "source": [
    "### 2. Get Cloud Formation stack output variables\n",
    "\n",
    "We also need to grab some key values from the infrastructure we provisioned using CloudFormation. To do this, we will list the outputs from the stack and store this in \"outputs\" to be used later.\n",
    "\n",
    "You can ignore any \"PythonDeprecationWarning\" warnings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f5ce6c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import boto3\n",
    "\n",
    "cfn = boto3.client('cloudformation')\n",
    "\n",
    "def get_cfn_outputs(stackname):\n",
    "    outputs = {}\n",
    "    for output in cfn.describe_stacks(StackName=stackname)['Stacks'][0]['Outputs']:\n",
    "        outputs[output['OutputKey']] = output['OutputValue']\n",
    "    return outputs\n",
    "\n",
    "## Setup variables to use for the rest of the demo\n",
    "cloudformation_stack_name = \"semantic-search\"\n",
    "\n",
    "outputs = get_cfn_outputs(cloudformation_stack_name)\n",
    "\n",
    "bucket = outputs['s3BucketTraining']\n",
    "aos_host = outputs['OpenSearchDomainEndpoint']\n",
    "\n",
    "outputs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bdbf3a68",
   "metadata": {},
   "source": [
    "### 3. Prepare training dataset\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "862612d3",
   "metadata": {},
   "source": [
    "Use [sentence transformer](https://www.sbert.net/docs/training/overview.html) to train the model. The input of training data is:\n",
    "* pair of text/sentence\n",
    "* label to indicate the semantic similarity between the sentences\n",
    "\n",
    "For demo purpose, we set orginal \"question\" and \"answer\" pair semantic similarity to \"1.0\". The \"question\" and \"answer in previous row\" semantic similarity to \"0.0\". In production, we shall use our training dataset based on business requirements.\n",
    "\n",
    "To reduce data load time, we only load 10,000 rows data for training.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a641b7c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import pandas as pd\n",
    "\n",
    "def load_pqa(file_name,number_rows=10000):\n",
    "    qa_list = []\n",
    "    df = pd.DataFrame(columns=('question', 'answer','label'))\n",
    "    with open(file_name) as f:\n",
    "        i=0\n",
    "        previous_row_data = None\n",
    "        for line in f:\n",
    "            data = json.loads(line)\n",
    "            df.loc[i] = [data['question_text'],data['answers'][0]['answer_text'],1.0]\n",
    "            i+=1\n",
    "            if previous_row_data is not None:\n",
    "                df.loc[i] = [data['question_text'],previous_row_data['answers'][0]['answer_text'],0.0]\n",
    "            previous_row_data = data\n",
    "            i+=1\n",
    "            if(i == number_rows*2):\n",
    "                break\n",
    "    return df\n",
    "\n",
    "\n",
    "qa_list = load_pqa('amazon-pqa/amazon_pqa_headsets.json')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dd2b5c30",
   "metadata": {},
   "source": [
    "We split the dataset into training dataset, validation dataset and test dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eab43a28",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import DataLoader\n",
    "import math\n",
    "from sentence_transformers import SentenceTransformer,  LoggingHandler, losses, models, util\n",
    "from sentence_transformers.evaluation import EmbeddingSimilarityEvaluator\n",
    "from sentence_transformers.readers import InputExample\n",
    "import logging\n",
    "from datetime import datetime\n",
    "import sys\n",
    "import os\n",
    "import gzip\n",
    "import csv\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sentence_transformers.readers import InputExample\n",
    "\n",
    "train_set,test_set = train_test_split(qa_list,test_size=0.2,shuffle=True)\n",
    "training_set, validation_set = train_test_split(train_set,test_size=0.2)\n",
    "\n",
    "def create_input_sample(data_set):\n",
    "    train_samples = []\n",
    "    for index,row in data_set.iterrows():\n",
    "        input_example = InputExample(texts=[row['question'], row['answer']], label=row['label'])\n",
    "        train_samples.append(input_example)\n",
    "    return train_samples\n",
    "\n",
    "training_samples = create_input_sample(training_set)\n",
    "validation_samples = create_input_sample(validation_set)\n",
    "test_samples = create_input_sample(test_set)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1fce0a6",
   "metadata": {},
   "source": [
    "### 6. Train the model\n",
    "\n",
    "#### Note: It will take more than 20 minutes to complete the training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e07a34a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import DataLoader\n",
    "from sentence_transformers import SentenceTransformer,  LoggingHandler, losses, models, util\n",
    "from sentence_transformers.evaluation import EmbeddingSimilarityEvaluator\n",
    "\n",
    "model_name = \"sentence-transformers/distilbert-base-nli-stsb-mean-tokens\"\n",
    "train_batch_size = 16\n",
    "num_epochs = 1\n",
    "model_save_path = 'output/fine_tuned_'+model_name.replace(\"/\", \"-\")\n",
    "\n",
    "word_embedding_model = models.Transformer(model_name)\n",
    "\n",
    "pooling_model = models.Pooling(word_embedding_model.get_word_embedding_dimension(),\n",
    "                               pooling_mode_mean_tokens=True,\n",
    "                               pooling_mode_cls_token=False,\n",
    "                               pooling_mode_max_tokens=False)\n",
    "\n",
    "model = SentenceTransformer(modules=[word_embedding_model, pooling_model])\n",
    "\n",
    "train_dataloader = DataLoader(training_samples, shuffle=True, batch_size=train_batch_size)\n",
    "train_loss = losses.CosineSimilarityLoss(model=model)\n",
    "\n",
    "evaluator = EmbeddingSimilarityEvaluator.from_input_examples(validation_samples, name='pqa-valucation')\n",
    "\n",
    "\n",
    "warmup_steps = math.ceil(len(train_dataloader) * num_epochs  * 0.1) #10% of train data for warm-up\n",
    "\n",
    "\n",
    "\n",
    "# Train the model\n",
    "model.fit(train_objectives=[(train_dataloader, train_loss)],\n",
    "          evaluator=evaluator,\n",
    "          epochs=num_epochs,\n",
    "          evaluation_steps=1000,\n",
    "          warmup_steps=warmup_steps,\n",
    "          output_path=model_save_path)\n",
    "\n",
    "\n",
    "model = SentenceTransformer(model_save_path)\n",
    "test_evaluator = EmbeddingSimilarityEvaluator.from_input_examples(test_samples, name='pqa-test')\n",
    "test_evaluator(model, output_path=model_save_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8fdf861b",
   "metadata": {},
   "source": [
    "### 7. Upload the fine-tuned model into S3"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "261c5f0e",
   "metadata": {},
   "source": [
    "Create tar file for the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e33018b",
   "metadata": {},
   "outputs": [],
   "source": [
    "!cd output/fine_tuned_sentence-transformers-distilbert-base-nli-stsb-mean-tokens && tar czvf ../model.tar.gz *"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "efc29e97",
   "metadata": {},
   "source": [
    "Upload the model to S3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f055c625",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sagemaker\n",
    "\n",
    "sagemaker_session = sagemaker.Session()\n",
    "role = sagemaker.get_execution_role()\n",
    "\n",
    "inputs = sagemaker_session.upload_data(path='output/model.tar.gz', key_prefix='fine-tuned-transformers-model')\n",
    "inputs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "61e86f19",
   "metadata": {},
   "source": [
    "### 8. Deploy the BERT model to SageMaker Endpoint\n",
    "\n",
    "First we need to create a PyTorchModel object. The deploy() method on the model object creates an endpoint which serves prediction requests in real-time. If the instance_type is set to a SageMaker instance type (e.g. ml.m5.large) then the model will be deployed on SageMaker. If the instance_type parameter is set to local then it will be deployed locally as a Docker container and ready for testing locally.\n",
    "\n",
    "First we need to create a Predictor class to accept TEXT as input and output JSON. The default behaviour is to accept a numpy array.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0151a880",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.pytorch import PyTorch, PyTorchModel\n",
    "from sagemaker.predictor import Predictor\n",
    "from sagemaker import get_execution_role\n",
    "\n",
    "class StringPredictor(Predictor):\n",
    "    def __init__(self, endpoint_name, sagemaker_session):\n",
    "        super(StringPredictor, self).__init__(endpoint_name, sagemaker_session, content_type='text/plain')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "376fc16a",
   "metadata": {},
   "source": [
    "Deploy the BERT model to Sagemaker Endpoint\n",
    "\n",
    "#### Note: This process will take serveral minutes to complete."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96d6bf1c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "pytorch_model = PyTorchModel(model_data = inputs, \n",
    "                             role=role, \n",
    "                             entry_point ='inference.py',\n",
    "                             source_dir = './code',\n",
    "                             py_version = 'py39', \n",
    "                             framework_version = '1.13.1',\n",
    "                             predictor_cls=StringPredictor)\n",
    "\n",
    "predictor = pytorch_model.deploy(instance_type='ml.m5d.large', \n",
    "                                 initial_instance_count=1, \n",
    "                                 endpoint_name = f'semantic-search-model-{int(time.time())}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4f780ad7",
   "metadata": {},
   "source": [
    "### 9. Test the SageMaker Endpoint.\n",
    "\n",
    "Now SageMaker Endpoind is ready, we can test the endpoind. Same as before, input is text data, output is vector. The difference is the BERT model is fune tuned with our business domain data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76a5efe1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "original_payload = 'Does this work with xbox?'\n",
    "features = predictor.predict(original_payload)\n",
    "vector_data = json.loads(features)\n",
    "\n",
    "vector_data\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1281a98d",
   "metadata": {},
   "source": [
    "### 10. Create an Amazon OpenSearch Service cluster connection.\n",
    "Next, we'll use Python API to set up connection with Amazon OpenSearch Service 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": "markdown",
   "id": "3a1b6c4c",
   "metadata": {},
   "source": [
    "Use Python API to set up connection with Amazon OpenSearch Cluster"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "416b0004",
   "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": "95ee25e8",
   "metadata": {},
   "source": [
    "### 11. 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": "e08e279f",
   "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": "352ec669",
   "metadata": {},
   "source": [
    "As we have created index in previous lab, let's delete the index. If this is the first time you're running this, you can skip this step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a6ed6d67",
   "metadata": {},
   "outputs": [],
   "source": [
    "aos_client.indices.delete(index=\"nlp_pqa\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cbd577fb",
   "metadata": {},
   "source": [
    "Using the above index definition, we now need to create the index in Amazon OpenSearch Service"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da3c6ea3",
   "metadata": {},
   "outputs": [],
   "source": [
    "aos_client.indices.create(index=\"nlp_pqa\",body=knn_index,ignore=400)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e888fba6",
   "metadata": {},
   "source": [
    "Let's verify the created index information"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "623a90e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "aos_client.indices.get(index=\"nlp_pqa\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "caf7e4de",
   "metadata": {},
   "source": [
    "### 12. 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": "9f093cd2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "from tqdm.contrib.concurrent import process_map\n",
    "from multiprocessing import cpu_count\n",
    "\n",
    "\n",
    "def load_pqa_as_json(file_name,number_rows=1000):\n",
    "    result=[]\n",
    "    with open(file_name) as f:\n",
    "        i=0\n",
    "        for line in f:\n",
    "            data = json.loads(line)\n",
    "            result.append(data)\n",
    "            i+=1\n",
    "            if(i == number_rows):\n",
    "                break\n",
    "    return result\n",
    "\n",
    "\n",
    "qa_list_json = load_pqa_as_json('amazon-pqa/amazon_pqa_headsets.json',number_rows=1000)\n",
    "\n",
    "\n",
    "def es_import(question):\n",
    "    vector = json.loads(predictor.predict(question[\"question_text\"]))\n",
    "    aos_client.index(index='nlp_pqa',\n",
    "             body={\"question_vector\": vector, \"question\": question[\"question_text\"],\"answer\":question[\"answers\"][0][\"answer_text\"]}\n",
    "            )\n",
    "        \n",
    "workers = 4 * cpu_count()\n",
    "    \n",
    "process_map(es_import, qa_list_json, max_workers=workers,chunksize=1000)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ecd4119",
   "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": "7d9efee1",
   "metadata": {},
   "outputs": [],
   "source": [
    "res = aos_client.search(index=\"nlp_pqa\", body={\"query\": {\"match_all\": {}}})\n",
    "print(\"Got %d Hits:\" % res['hits']['total']['value'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "81fd4d06",
   "metadata": {},
   "source": [
    "### 13. Generate vector for user input query \n",
    "\n",
    "Next, we'll use SageMaker Endpoint to convert our input question \"does this work with xbox?\" into a vector. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa0e19ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_raw_sentences = ['does this work with xbox?']\n",
    "client = boto3.client('sagemaker-runtime')\n",
    "ENDPOINT_NAME = predictor.endpoint\n",
    "response = client.invoke_endpoint(EndpointName=ENDPOINT_NAME,\n",
    "                                       ContentType='text/plain',\n",
    "                                       Body=query_raw_sentences[0])\n",
    "\n",
    "search_vector = json.loads((response['Body'].read()))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "95a38dba",
   "metadata": {},
   "source": [
    "### 14. Search vector with \"Semantic Search\" \n",
    "\n",
    "Now that we have vector in Amazon OpenSearch Service and a vector for our query question, let's perform a KNN search in Amazon OpenSearch.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89937297",
   "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": "2f595c05",
   "metadata": {},
   "source": [
    "### 15. 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": "d4b6d129",
   "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)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae3958bf",
   "metadata": {},
   "source": [
    "### 16. 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)."
   ]
  }
 ],
 "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}