Module 6 - Lab2 - Sematic Search with Neural Search Remote Model.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e87dc259",
   "metadata": {},
   "source": [
    "# Semantic Search with OpenSearch Neural Search "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b0cfd51d",
   "metadata": {},
   "source": [
    "We will use Neural Search plugin in OpenSearch to implement 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 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": "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": "4a3fa4b0",
   "metadata": {},
   "source": [
    "### 3. Install OpenSearch ML Python library"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58a1c491",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install opensearch-py-ml\n",
    "!pip install accelerate\n",
    "!pip install deprecated\n",
    "!pip install requests-aws4auth"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "05c00375",
   "metadata": {},
   "source": [
    "Now we need to restart the kernel by running below cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f94df946",
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import display_html\n",
    "def restartkernel() :\n",
    "    display_html(\"<script>Jupyter.notebook.kernel.restart()</script>\",raw=True)\n",
    "restartkernel()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0aa614bc",
   "metadata": {},
   "source": [
    "### 4. Import library\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1688f4e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import boto3\n",
    "import re\n",
    "import time"
   ]
  },
  {
   "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. \n",
    "\n",
    "Before we can run any queries, we need to download the Amazon Product Question and Answer data from : https://registry.opendata.aws/amazon-pqa/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58fca957",
   "metadata": {},
   "outputs": [],
   "source": [
    "!aws s3 cp --no-sign-request s3://amazon-pqa/amazon_pqa_headsets.json ./amazon-pqa/amazon_pqa_headsets.json"
   ]
  },
  {
   "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": "2f54a349",
   "metadata": {},
   "source": [
    "### 6. 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.\n",
    "\n",
    "#### 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": "81dc45a2",
   "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": "code",
   "execution_count": null,
   "id": "0cee3dd0",
   "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",
    "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": "62421acd",
   "metadata": {},
   "source": [
    "### 7. Register Model Groups and External Connector "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "82210464",
   "metadata": {},
   "source": [
    "Initialize OpenSearch authentication with AWS auth."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bda0e264",
   "metadata": {},
   "outputs": [],
   "source": [
    "import boto3\n",
    "import requests \n",
    "from requests_aws4auth import AWS4Auth\n",
    "\n",
    "service = 'es'\n",
    "credentials = boto3.Session().get_credentials()\n",
    "awsauth = AWS4Auth(credentials.access_key, credentials.secret_key, region, service, session_token=credentials.token)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8abf8665",
   "metadata": {},
   "source": [
    "Register model group\n",
    "\n",
    "---\n",
    "### Note: Before run this cell, make sure you have completed the steps in the lab instruction \"Map the ML role in OpenSearch Dashboards\". \n",
    "\n",
    "If you don't complete the lab instruction steps, you will get \"403\" error, error message likes \"There is error in creating connector{\"error\":{\"root_cause\":[{\"type\":\"security_exception\",\"reason\":\"no permissions for ...\". \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76f13ba2",
   "metadata": {},
   "outputs": [],
   "source": [
    "path = '_plugins/_ml/model_groups/_register'\n",
    "url = 'https://' + aos_host + '/' + path\n",
    "\n",
    "payload = {\n",
    "  \"name\": \"remote_model_groups_for_embedding\",\n",
    "  \"description\": \"A model group for remote models\"\n",
    "}\n",
    "\n",
    "headers = {\"Content-Type\": \"application/json\"}\n",
    "\n",
    "r = requests.post(url, auth=awsauth, json=payload, headers=headers)\n",
    "print(r.status_code)\n",
    "if r.status_code == 200:\n",
    "    data = json.loads(r.text)\n",
    "\n",
    "    model_group_id = data['model_group_id']\n",
    "    print(\"model group id:\" + model_group_id)\n",
    "else:\n",
    "    raise Exception(\"There is error in creating model groups\" + str(r.text))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "02e661b9",
   "metadata": {},
   "source": [
    "Uncomment the following code if you want delete model group. Replace  `{model_group_id}` with the value you want delete."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a74f72b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# path = '_plugins/_ml/model_groups/{model_group_id}'\n",
    "# url = 'https://' + aos_host + '/' + path\n",
    "\n",
    "\n",
    "# headers = {\"Content-Type\": \"application/json\"}\n",
    "\n",
    "# r = requests.delete(url, auth=awsauth, headers=headers)\n",
    "# print(r.text)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "64ef5d5d",
   "metadata": {},
   "source": [
    "Get account id"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "636971f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "account_id = boto3.client('sts').get_caller_identity().get('Account')\n",
    "print(account_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3136c308",
   "metadata": {},
   "source": [
    "### Register SageMaker Connector\n",
    "\n",
    "\n",
    "---\n",
    "\n",
    "There are several important parameters to register SageMaker connector. Please refer [OpenSearch Connector blueprints](https://opensearch.org/docs/latest/ml-commons-plugin/extensibility/blueprints/#configuration-options) for more information.\n",
    "\n",
    "* pre_process_function: Function used to pre-process data before sending to embedding model. Example script used to preprocess the input data is as following:\n",
    "```\n",
    "    StringBuilder builder = new StringBuilder();\n",
    "    builder.append(\"\\\"\");\n",
    "    builder.append(params.text_docs[0]);\n",
    "    builder.append(\"\\\"\");\n",
    "    def parameters = \"{\" +\"\\\"inputs\\\":\" + builder + \"}\";\n",
    "    return \"{\" +\"\\\"parameters\\\":\" + parameters + \"}\";\n",
    "```\n",
    "\n",
    "* request_body: Define the data structure sent to enmbedding model. Example input data to GPT-J embedding model is like following:\n",
    "\n",
    "```\n",
    "{ \"text_inputs\": \"${parameters.inputs}\"}\n",
    "\n",
    "```\n",
    "\n",
    "* post_process_function: Function used to post-process data after getting embedding data from ML model. Example script used to post-process the model output data:\n",
    "```\n",
    "    def name = \"sentence_embedding\";\n",
    "    def dataType = \"FLOAT32\";\n",
    "    if (params.embedding == null || params.embedding.length == 0) {\n",
    "        return null;\n",
    "    }\n",
    "    def shape = [params.embedding[0].length];\n",
    "    def json = \"{\" +\n",
    "               \"\\\"name\\\":\\\"\" + name + \"\\\",\" +\n",
    "               \"\\\"data_type\\\":\\\"\" + dataType + \"\\\",\" +\n",
    "               \"\\\"shape\\\":\" + shape + \",\" +\n",
    "               \"\\\"data\\\":\" + params.embedding[0] +\n",
    "               \"}\";\n",
    "    return json;\n",
    "```\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8644b9f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "path = '_plugins/_ml/connectors/_create'\n",
    "url = 'https://' + aos_host + '/' + path\n",
    "\n",
    "sagemaker_url = \"https://runtime.sagemaker.us-east-1.amazonaws.com/endpoints/\" + outputs[\"EmbeddingEndpointName\"]+ \"/invocations\"\n",
    "role_arn = \"arn:aws:iam::\" + account_id + \":role/opensearch-sagemaker-role\"\n",
    "\n",
    "payload = {\n",
    "   \"name\": \"sagemaker embedding\",\n",
    "   \"description\": \"Remote connector for Sagemaker embedding model\",\n",
    "   \"version\": 1,\n",
    "   \"protocol\": \"aws_sigv4\",\n",
    "   \"credential\": {\n",
    "      \"roleArn\": role_arn\n",
    "   },\n",
    "   \"parameters\": {\n",
    "      \"region\": \"us-east-1\",\n",
    "      \"service_name\": \"sagemaker\"\n",
    "   },\n",
    "   \"actions\": [\n",
    "      {\n",
    "         \"action_type\": \"predict\",\n",
    "         \"method\": \"POST\",\n",
    "         \"headers\": {\n",
    "            \"content-type\": \"application/json\"\n",
    "         },\n",
    "         \"url\": sagemaker_url,\n",
    "         \"pre_process_function\": '\\n    StringBuilder builder = new StringBuilder();\\n    builder.append(\"\\\\\"\");\\n    builder.append(params.text_docs[0]);\\n    builder.append(\"\\\\\"\");\\n    def parameters = \"{\" +\"\\\\\"inputs\\\\\":\" + builder + \"}\";\\n    return \"{\" +\"\\\\\"parameters\\\\\":\" + parameters + \"}\";\\n    ', \n",
    "         \"request_body\": \"{ \\\"text_inputs\\\": \\\"${parameters.inputs}\\\"}\",\n",
    "         \"post_process_function\": '\\n    def name = \"sentence_embedding\";\\n    def dataType = \"FLOAT32\";\\n    if (params.embedding == null || params.embedding.length == 0) {\\n        return null;\\n    }\\n    def shape = [params.embedding[0].length];\\n    def json = \"{\" +\\n               \"\\\\\"name\\\\\":\\\\\"\" + name + \"\\\\\",\" +\\n               \"\\\\\"data_type\\\\\":\\\\\"\" + dataType + \"\\\\\",\" +\\n               \"\\\\\"shape\\\\\":\" + shape + \",\" +\\n               \"\\\\\"data\\\\\":\" + params.embedding[0] +\\n               \"}\";\\n    return json;\\n    '\n",
    "\n",
    "      }\n",
    "   ]\n",
    "}\n",
    "headers = {\"Content-Type\": \"application/json\"}\n",
    "\n",
    "r = requests.post(url, auth=awsauth, json=payload, headers=headers)\n",
    "print(r.status_code)\n",
    "if r.status_code == 200:\n",
    "    data = json.loads(r.text)\n",
    "\n",
    "    sagemaker_connector_id = data['connector_id']\n",
    "    print(\"SageMaker connector id:\" + sagemaker_connector_id)\n",
    "else:\n",
    "    raise Exception(\"There is error in creating connector\" + str(r.text))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "308470fa",
   "metadata": {},
   "source": [
    "### Register Bedrock Connector\n",
    "\n",
    "---\n",
    "\n",
    "Register Bedrock connecter is similiar with SageMaker connector, only with some parameters difference. Please refere [Amazon Bedrock connector](https://opensearch.org/docs/latest/ml-commons-plugin/extensibility/connectors/#amazon-bedrock-connector) on how to register Bedrock connector.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8db0200",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "path = '_plugins/_ml/connectors/_create'\n",
    "url = 'https://' + aos_host + '/' + path\n",
    "\n",
    "bedrock_url = \"https://bedrock-runtime.\" + region + \".amazonaws.com/model/amazon.titan-embed-text-v1/invoke\"\n",
    "role_arn = \"arn:aws:iam::\" + account_id + \":role/opensearch-sagemaker-role\"\n",
    "\n",
    "payload = {\n",
    "  \"name\": \"Amazon Bedrock Connector: embedding\",\n",
    "  \"description\": \"The connector to bedrock Titan embedding model\",\n",
    "  \"version\": 1,\n",
    "  \"protocol\": \"aws_sigv4\",\n",
    "  \"parameters\": {\n",
    "    \"region\": region,\n",
    "    \"service_name\": \"bedrock\"\n",
    "  },\n",
    "  \"credential\": {\n",
    "    \"roleArn\": role_arn\n",
    "  },\n",
    "  \"actions\": [\n",
    "    {\n",
    "      \"action_type\": \"predict\",\n",
    "      \"method\": \"POST\",\n",
    "      \"url\": bedrock_url,\n",
    "      \"headers\": {\n",
    "        \"content-type\": \"application/json\",\n",
    "        \"x-amz-content-sha256\": \"required\"\n",
    "      },\n",
    "      \"pre_process_function\": \"\\n    StringBuilder builder = new StringBuilder();\\n    builder.append(\\\"\\\\\\\"\\\");\\n    String first = params.text_docs[0];\\n    builder.append(first);\\n    builder.append(\\\"\\\\\\\"\\\");\\n    def parameters = \\\"{\\\" +\\\"\\\\\\\"inputText\\\\\\\":\\\" + builder + \\\"}\\\";\\n    return  \\\"{\\\" +\\\"\\\\\\\"parameters\\\\\\\":\\\" + parameters + \\\"}\\\";\",\n",
    "      \"request_body\": \"{ \\\"inputText\\\": \\\"${parameters.inputText}\\\" }\",\n",
    "      \"post_process_function\": \"\\n      def name = \\\"sentence_embedding\\\";\\n      def dataType = \\\"FLOAT32\\\";\\n      if (params.embedding == null || params.embedding.length == 0) {\\n        return params.message;\\n      }\\n      def shape = [params.embedding.length];\\n      def json = \\\"{\\\" +\\n                 \\\"\\\\\\\"name\\\\\\\":\\\\\\\"\\\" + name + \\\"\\\\\\\",\\\" +\\n                 \\\"\\\\\\\"data_type\\\\\\\":\\\\\\\"\\\" + dataType + \\\"\\\\\\\",\\\" +\\n                 \\\"\\\\\\\"shape\\\\\\\":\\\" + shape + \\\",\\\" +\\n                 \\\"\\\\\\\"data\\\\\\\":\\\" + params.embedding +\\n                 \\\"}\\\";\\n      return json;\\n    \"\n",
    "    }\n",
    "  ]\n",
    "}\n",
    "headers = {\"Content-Type\": \"application/json\"}\n",
    "\n",
    "r = requests.post(url, auth=awsauth, json=payload, headers=headers)\n",
    "print(r.status_code)\n",
    "if r.status_code == 200:\n",
    "    data = json.loads(r.text)\n",
    "\n",
    "    bedrock_connector_id = data['connector_id']\n",
    "    print(\"Bedrock connector id:\" + bedrock_connector_id)\n",
    "else:\n",
    "    raise Exception(\"There is error in creating connector\" + str(r.text))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b4de03e",
   "metadata": {},
   "source": [
    "### Select Embedding Model\n",
    "\n",
    "Here you can choose SageMaker or Bedrock for embedding. After the Dropbox is show, select one model you will use for the following steps.\n",
    "\n",
    "Depends on how you create this lab environment, you may have Bedrock in your environment or not. If you don't have Bedorck in your environment, you have to set `is_bedrock_available` to False.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ea33a27",
   "metadata": {},
   "outputs": [],
   "source": [
    "is_bedrock_available=True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6adc9488",
   "metadata": {},
   "outputs": [],
   "source": [
    "from ipywidgets import Dropdown\n",
    "\n",
    "llm_selection = [\n",
    "    \"SageMaker\",\n",
    "    \"Bedrock\",\n",
    "]\n",
    "\n",
    "llm_dropdown = Dropdown(\n",
    "    options=llm_selection,\n",
    "    value=\"SageMaker\",\n",
    "    description=\"Select embedding model\",\n",
    "    style={\"description_width\": \"initial\"},\n",
    "    layout={\"width\": \"max-content\"},\n",
    ")\n",
    "display(llm_dropdown)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c2904cf6",
   "metadata": {},
   "outputs": [],
   "source": [
    "llm_category = llm_dropdown.value\n",
    "\n",
    "if not is_bedrock_available:\n",
    "    llm_category = \"SageMaker\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "39c89a5a",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"You selected {0} as embedding model\".format(llm_category))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01630c90",
   "metadata": {},
   "outputs": [],
   "source": [
    "match llm_category:\n",
    "    case \"SageMaker\":\n",
    "        connector_id = sagemaker_connector_id\n",
    "    case \"Bedrock\":\n",
    "        connector_id = bedrock_connector_id\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b6efd116",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"connector id: \" + connector_id)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a19d7ccd",
   "metadata": {},
   "source": [
    "### 8. Register  Remote Model\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e93193ae",
   "metadata": {},
   "source": [
    "Register a remote model with connector and model groups"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c77f499",
   "metadata": {},
   "outputs": [],
   "source": [
    "path = '_plugins/_ml/models/_register'\n",
    "url = 'https://' + aos_host + '/' + path\n",
    "\n",
    "sagemaker_payload = {\n",
    "    \"name\": \"sagemaker-opensearch-embedding\",\n",
    "    \"function_name\": \"remote\",\n",
    "    \"model_group_id\": model_group_id,\n",
    "    \"description\": \"opensearch gpt-j embedding\",\n",
    "    \"connector_id\": connector_id\n",
    "}\n",
    "\n",
    "bedrock_payload = {\n",
    "    \"name\": \"bedrock-opensearch-embedding\",\n",
    "    \"function_name\": \"remote\",\n",
    "    \"model_group_id\": model_group_id,\n",
    "    \"description\": \"bedrock titan embedding\",\n",
    "    \"connector_id\": connector_id\n",
    "}\n",
    "\n",
    "match llm_category:\n",
    "    case \"SageMaker\":\n",
    "        payload = sagemaker_payload\n",
    "    case \"Bedrock\":\n",
    "        payload = bedrock_payload\n",
    "        \n",
    "headers = {\"Content-Type\": \"application/json\"}\n",
    "\n",
    "r = requests.post(url, auth=awsauth, json=payload, headers=headers)\n",
    "print(r.status_code)\n",
    "if r.status_code == 200:\n",
    "    data = json.loads(r.text)\n",
    "\n",
    "    model_id = data['model_id']\n",
    "    print(model_id)\n",
    "else:\n",
    "    raise Exception(\"There is error in registering model\" + str(r.text))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5dc9d104",
   "metadata": {},
   "source": [
    "### 9. Load the model for inference."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52cba357",
   "metadata": {},
   "outputs": [],
   "source": [
    "from opensearch_py_ml.ml_commons import MLCommonClient\n",
    "\n",
    "ml_client = MLCommonClient(aos_client)\n",
    "\n",
    "load_model_output = ml_client.deploy_model(model_id)\n",
    "\n",
    "print(load_model_output)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c54be9c",
   "metadata": {},
   "source": [
    "### 10.Get the model detailed information and test the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1211a76d",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_info = ml_client.get_model_info(model_id)\n",
    "\n",
    "print(model_info)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1e6b673a",
   "metadata": {},
   "source": [
    "Test the remote embedding model "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "048bc162",
   "metadata": {},
   "outputs": [],
   "source": [
    "path = '_plugins/_ml/models/' + model_id + '/_predict'\n",
    "url = 'https://' + aos_host + '/' + path\n",
    "\n",
    "\n",
    "sagemaker_payload = {\n",
    "  \"parameters\": {\n",
    "    \"inputs\": \"this is test\"\n",
    "  }\n",
    "}\n",
    "\n",
    "bedrock_payload = {\n",
    "  \"parameters\": {\n",
    "    \"inputText\": \"this is test\"\n",
    "  }\n",
    "}\n",
    "\n",
    "match llm_category:\n",
    "    case \"SageMaker\":\n",
    "        payload = sagemaker_payload\n",
    "    case \"Bedrock\":\n",
    "        payload = bedrock_payload\n",
    "        \n",
    "headers = {\"Content-Type\": \"application/json\"}\n",
    "\n",
    "r = requests.post(url, auth=awsauth, json=payload, headers=headers)\n",
    "print(r.status_code)\n",
    "if r.status_code == 200:\n",
    "    data = json.loads(r.text)\n",
    "    print(data)\n",
    "else:\n",
    "    raise Exception(\"There is error in calling model\" + str(r.text))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3625b5cf",
   "metadata": {},
   "source": [
    "### 11. Create pipeline to convert text into vector with BERT model\n",
    "We will use the just uploaded model to convert `qestion` field into vector(embedding) and stored into `question_vector` field."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc810643",
   "metadata": {},
   "outputs": [],
   "source": [
    "pipeline={\n",
    "  \"description\": \"An semantic search pipeline\",\n",
    "  \"processors\" : [\n",
    "    {\n",
    "      \"text_embedding\": {\n",
    "        \"model_id\": model_id,\n",
    "        \"field_map\": {\n",
    "           \"question\": \"question_vector\"\n",
    "        }\n",
    "      }\n",
    "    }\n",
    "  ]\n",
    "}\n",
    "pipeline_id = 'nlp_pipeline'\n",
    "aos_client.ingest.put_pipeline(id=pipeline_id,body=pipeline)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c431a804",
   "metadata": {},
   "source": [
    "Verify pipeline is created succefuflly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13ff2f91",
   "metadata": {},
   "outputs": [],
   "source": [
    "aos_client.ingest.get_pipeline(id=pipeline_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "beaabc1e",
   "metadata": {},
   "source": [
    "### 12. 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.\n",
    "\n",
    "Here we need to define different index for SageMaker and Bedrock embedding because the vector dimension is different. SageMaker embedding dimension is 4096, Bedrock embedding dimension is 1536."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5eba5754",
   "metadata": {},
   "outputs": [],
   "source": [
    "sagemaker_knn_index = {\n",
    "    \"settings\": {\n",
    "        \"index.knn\": True,\n",
    "        \"index.knn.space_type\": \"cosinesimil\",\n",
    "        \"default_pipeline\": pipeline_id,\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\": 4096,\n",
    "                \"method\": {\n",
    "                    \"name\": \"hnsw\",\n",
    "                    \"space_type\": \"l2\",\n",
    "                    \"engine\": \"faiss\"\n",
    "                },\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",
    "\n",
    "bedrock_knn_index = {\n",
    "    \"settings\": {\n",
    "        \"index.knn\": True,\n",
    "        \"index.knn.space_type\": \"cosinesimil\",\n",
    "        \"default_pipeline\": pipeline_id,\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\": 1536,\n",
    "                \"method\": {\n",
    "                    \"name\": \"hnsw\",\n",
    "                    \"space_type\": \"l2\",\n",
    "                    \"engine\": \"faiss\"\n",
    "                },\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",
    "\n",
    "match llm_category:\n",
    "    case \"SageMaker\":\n",
    "        knn_index = sagemaker_knn_index\n",
    "    case \"Bedrock\":\n",
    "        knn_index = bedrock_knn_index\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1330502a",
   "metadata": {},
   "source": [
    "## Note: If this is the first time you're running this, you can comment this line code which try to delete index created before."
   ]
  },
  {
   "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": [
    "print(knn_index)\n",
    "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": [
    "### 13. Load the raw data into the Index\n",
    "Next, let's load the headset enhanced PQA data into the index we've just created. During ingest data, `question` field will also be converted to vector(embedding) by the `nlp_pipeline` we defined."
   ]
  },
  {
   "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",
    "    answer=qa_list[\"answer\"][i]\n",
    "    i+=1\n",
    "    aos_client.index(index='nlp_pqa',body={\"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'])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "de9b827c",
   "metadata": {},
   "source": [
    "### 14. Search vector with \"Semantic Search\" \n",
    "\n",
    "We can search the data with neural search.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c5f4e81",
   "metadata": {},
   "outputs": [],
   "source": [
    "query={\n",
    "  \"_source\": {\n",
    "        \"exclude\": [ \"question_vector\" ]\n",
    "    },\n",
    "  \"size\": 30,\n",
    "  \"query\": {\n",
    "    \"neural\": {\n",
    "      \"question_vector\": {\n",
    "        \"query_text\": \"does this work with xbox?\",\n",
    "        \"model_id\": model_id,\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['_source']['question'],hit['_source']['answer']]\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": [
    "### 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": "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": [
    "### 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).\n",
    "\n",
    "You can also use \"Compare search results\" in Search relevance plugin to compare search relevance side by side. Please refer the lab \"Option 2: OpenSearch Dashboard Dev Tools\" to compare search results."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dd1bd18f",
   "metadata": {},
   "source": [
    "### 17. Choose another embedding model \n",
    "\n",
    "Go back to [Select Embedding Model](#Select-Embedding-Model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c8d80ff",
   "metadata": {},
   "source": [
    "### 18. Summary\n",
    "With OpenSearch Neural Search remote model, embedding is automatically generated with model hosted in SageMaker or Bedrock. We don't need care about inference pipeline anymore. It makes the semantic search solution simple to develop and maintain. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6c15b678",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}