セマンティックリランキング (Amazon SageMaker 編)¶

from IPython.core.magic import register_cell_magic
from IPython import get_ipython
import ipywidgets as widgets

import boto3
import json
import time
import logging

from tqdm import tqdm
from datetime import datetime, timedelta
from functools import lru_cache

import pandas as pd
import numpy as np
from opensearchpy import OpenSearch, RequestsHttpConnection, AWSV4SignerAuth

import sagemaker
from sagemaker.huggingface import HuggingFaceModel, get_huggingface_llm_image_uri

try:
    sagemaker_region = default_region
    reranking_model_id_on_hf = "BAAI/bge-reranker-v2-m3"
    reranking_model_name = reranking_model_id_on_hf.lower().replace("/", "-")
    realtime_inference_endpoint_type = "instance"

    reranking_endpoint_name_prefix = f"{reranking_model_name}-{realtime_inference_endpoint_type}"
    response = search_sagemaker_inference_endpoint(reranking_endpoint_name_prefix, sagemaker_region)

    if len(response) == 0:
        print(f"Model {reranking_model_id_on_hf} is not deployed on a realtime {realtime_inference_endpoint_type} endpoint." )

        hub = {
          "HF_MODEL_ID": reranking_model_id_on_hf
        }
        
        role = sagemaker.get_execution_role()
        
        available_instance_types = list_available_instance_types_for_realtime_inference(instance_family="g5", region=sagemaker_region)
        if (len(available_instance_types)) > 0:
            instance_type = available_instance_types[0]
        else:
            print("There is no available gpu instance type. Trying to use cpu instance type.")
            available_instance_types = list_available_instance_types_for_realtime_inference(instance_family="c5", region=sagemaker_region)
            if (len(available_instance_types)) > 0:
                instance_type = available_instance_types[0]
            else:
                raise ValueException("There is no available instance types. Please change deployment option to on serverless endpoint, and try again.")
        print(f"Found an eligable instance type. {instance_type} will be used for a realtime {realtime_inference_endpoint_type} endpoint.")
        reranking_model = HuggingFaceModel(
            name=sagemaker.utils.name_from_base(reranking_model_name),
            env=hub, # configuration for loading model from Hub
            role=role, # iam role with permissions to create an Endpoint
            image_uri=get_huggingface_tei_image_uri(instance_type=instance_type, region=sagemaker_region)
        )
        print(f"start deploy {reranking_model_id_on_hf} on a realtime {realtime_inference_endpoint_type} endpoint.")
        reranking_model.deploy(
            endpoint_name=sagemaker.utils.name_from_base(reranking_endpoint_name_prefix),
            initial_instance_count=1,
            instance_type=instance_type
        )

        reranking_endpoint_name = reranking_model.endpoint_name

    else:
        print(f"Model {reranking_model_id_on_hf} is already deployed on a realtime {realtime_inference_endpoint_type} endpoint.")
        reranking_endpoint = response[0]["Endpoint"]
        reranking_endpoint_name = reranking_endpoint["EndpointName"]

    reranking_endpoint_url = f"https://runtime.sagemaker.{default_region}.amazonaws.com/endpoints/{reranking_endpoint_name}/invocations"
    print(f"\nendpoint name: {reranking_endpoint_name}")
    print(f"endpoint url: {reranking_endpoint_url}")
    
except Exception as e:
    print(e)

%%time

payload = {
    "query":"What is the capital city of America?",
    "texts":[
        "Carson City is the capital city of the American state of Nevada.",
        "The Commonwealth of the Northern Mariana Islands is a group of islands in the Pacific Ocean. Its capital is Saipan.",
        "Washington, D.C. (also known as simply Washington or D.C., and officially as the District of Columbia) is the capital of the United States. It is a federal district.",
        "Capital punishment (the death penalty) has existed in the United States since beforethe United States was a country. As of 2017, capital punishment is legal in 30 of the 50 states."
    ]
}

body = bytes(json.dumps(payload), 'utf-8')

sagemaker_runtime_client = boto3.client("sagemaker-runtime", region_name=sagemaker_region)
response = sagemaker_runtime_client.invoke_endpoint(
    EndpointName=reranking_endpoint_name,
    ContentType="application/json",
    Accept="application/json",
    Body=body
)

result = eval(response['Body'].read().decode('utf-8'))
print(json.dumps(result, indent=2))

print(json.dumps(sorted(result, key=lambda x: x['index']),indent=2))

cloudformation_stack_name = "search-lab-jp"
opensearch_cluster_endpoint = search_cloudformation_output(cloudformation_stack_name, "OpenSearchDomainEndpoint")

credentials = boto3.Session().get_credentials()
service_code = "es"
auth = AWSV4SignerAuth(credentials=credentials, region=default_region, service=service_code)
opensearch_client = OpenSearch(
    hosts=[{"host": opensearch_cluster_endpoint, "port": 443}],
    http_compress=True, 
    http_auth=auth,
    use_ssl=True,
    verify_certs=True,
    connection_class = RequestsHttpConnection
)
opensearch_client

opensearch_client.info()

print(f"model name: {reranking_model_name}")
print(f"endpoint name: {reranking_endpoint_name}")
print(f"endpoint url: {reranking_endpoint_url}")

cloudformation_stack_name = "search-lab-jp"
opensearch_connector_role_arn = search_cloudformation_output(cloudformation_stack_name, 'OpenSearchConnectorRoleArn')
opensearch_connector_role_arn

payload = {
  "name": reranking_model_name,
  "description": "Remote connector for "+ reranking_model_name,
  "version": 1,
  "protocol": "aws_sigv4",
  "credential": {
    "roleArn": opensearch_connector_role_arn
  },
  "parameters": {
    "region": default_region,
    "service_name": "sagemaker"
  },
  "actions": [
    {
      "action_type": "predict",
      "method": "POST",
      "url": reranking_endpoint_url,
      "headers": {
        "content-type": "application/json"
      },
      "pre_process_function": """
        def query_text = params.query_text;
        def text_docs = params.text_docs;
        def textDocsBuilder = new StringBuilder('[');
        for (int i=0; i<text_docs.length; i++) {
          textDocsBuilder.append('\"');
          textDocsBuilder.append(text_docs[i]);
          textDocsBuilder.append('\"');
          if (i<text_docs.length - 1) {
            textDocsBuilder.append(',');
          }
        }
        textDocsBuilder.append(']');
        def parameters = '{ \"query\": \"' + query_text + '\",  \"texts\": ' + textDocsBuilder.toString() + ' }';
        return  '{\"parameters\": ' + parameters + '}';
      """,
      "request_body": "{ \"query\": \"${parameters.query}\", \"texts\": ${parameters.texts} }",
      "post_process_function": """
        if (params.result == null) {
          throw new IllegalArgumentException("Post process function input is empty.");
        }
        def outputs = params.result;
        def scores = new Double[outputs.length];
        for (int i=0; i<outputs.length; i++) {
          def index = new BigDecimal(outputs[i].index.toString()).intValue();
          scores[index] = outputs[i].score;
        }
        def resultBuilder = new StringBuilder('[');
        for (int i=0; i<scores.length; i++) {
          resultBuilder.append(' {\"name\": \"similarity\", \"data_type\": \"FLOAT32\", \"shape\": [1],');
          resultBuilder.append('\"data\": [');
          resultBuilder.append(scores[i]);
          resultBuilder.append(']}');
          if (i<outputs.length - 1) {
            resultBuilder.append(',');
          }
        }
        resultBuilder.append(']');
        return resultBuilder.toString();
      """
    }
  ]
}

response = opensearch_client.http.post("/_plugins/_ml/connectors/_create", body=payload)
reranking_connector_id = response['connector_id']
print("reranking connector id: " + reranking_connector_id)

payload = {
    "name": reranking_model_name,
    "description": reranking_model_name,
    "function_name": "remote",
    "connector_id": reranking_connector_id
}
response = opensearch_client.http.post("/_plugins/_ml/models/_register?deploy=true", body=payload)

reranking_model_id = response['model_id']

for i in range(300):
    ml_model_status = opensearch_client.http.get("/_plugins/_ml/models/"+ reranking_model_id)
    model_state = ml_model_status.get("model_state")
    if model_state in ["DEPLOYED", "PARTIALLY_DEPLOYED"]:
        break
    time.sleep(1)

if model_state == "DEPLOYED":
    print("reranking model " + reranking_model_id + " is deployed successfully")
elif model_state == "PARTIALLY_DEPLOYED":
    print("reranking model " + reranking_model_id + " is deployed only partially")
else:
    raise Exception("reranking model " + reranking_model_id + " deployment failed")

print(ml_model_status)

payload = {
  "query_text": "What is the capital city of America?",
  "text_docs": [
    "Carson City is the capital city of the American state of Nevada.",
    "The Commonwealth of the Northern Mariana Islands is a group of islands in the Pacific Ocean. Its capital is Saipan.",
    "Washington, D.C. (also known as simply Washington or D.C., and officially as the District of Columbia) is the capital of the United States. It is a federal district.",
    "Capital punishment (the death penalty) has existed in the United States since beforethe United States was a country. As of 2017, capital punishment is legal in 30 of the 50 states."
  ]
}

opensearch_client.http.post("/_plugins/_ml/_predict/text_similarity/" + reranking_model_id, body=payload)

payload = {
  "parameters": {
    "query": "What is the capital city of America?",
    "texts": [
        "Carson City is the capital city of the American state of Nevada.",
        "The Commonwealth of the Northern Mariana Islands is a group of islands in the Pacific Ocean. Its capital is Saipan.",
        "Washington, D.C. (also known as simply Washington or D.C., and officially as the District of Columbia) is the capital of the United States. It is a federal district.",
        "Capital punishment (the death penalty) has existed in the United States since beforethe United States was a country. As of 2017, capital punishment is legal in 30 of the 50 states."
    ]
  }
}

result = opensearch_client.http.post("/_plugins/_ml/models/" + reranking_model_id + "/_predict", body=payload)
print(json.dumps(result, indent=2))

payload={
  "request_processors": request_processors,
  "phase_results_processors": phase_results_processors,
  "response_processors": [
    {
      "rerank": {
        "ml_opensearch": {
          "model_id": reranking_model_id
        },
        "context": {
          "document_fields": [
            "question"
          ]
        }
      }
    }
  ]
}
# パイプライン ID の指定
rerank_pipeline_id = "rerank_pipeline"
# パイプライン作成 API の呼び出し
response = opensearch_client.http.put("/_search/pipeline/" + rerank_pipeline_id, body=payload)
print(response)

response = opensearch_client.http.get("/_search/pipeline/" + rerank_pipeline_id)
print(response)

%%time
# search
index_name = "jsquad-neural-search"
query = "日本で梅雨がない場所は？"
payload = {
  "size": 30,
  "query": {
    "neural": {
      "context_embedding": {
        "query_text": query, # テキストをベクトルに変換し
        "k": 30 # クエリベクトルに近いベクトルのうち上位 5 つを返す
      }
    }
  },
  "_source" : False,
  "fields": ["question", "answers",  "context"]
}
# 検索 API を実行
response = opensearch_client.search(
    body = payload,
    index = index_name,
    filter_path = "hits.hits",
    search_pipeline = neural_search_pipeline_id #ベクトル変換を行うパイプラインを指定
)

# 結果を表示
pd.json_normalize(response["hits"]["hits"])

%%time
# search
index_name = "jsquad-neural-search"
query = "日本で梅雨がない場所は？"
payload = {
  "size": 30,
  "query": {
    "neural": {
      "context_embedding": {
        "query_text": query, # テキストをベクトルに変換し
        "k": 30 # クエリベクトルに近いベクトルのうち上位 3 つを返す
      }
    }
  },
  "ext": {
    "rerank": {
      "query_context": {
         "query_text_path": "query.neural.context_embedding.query_text"
      }
    }
  },
  "_source" : False,
  "fields": ["question", "answers",  "context"]
}
# 検索 API を実行
response = opensearch_client.search(
    body = payload,
    index = index_name,
    filter_path = "hits.hits",
    search_pipeline = rerank_pipeline_id #ベクトル変換を行うパイプラインを指定
)

# 結果を表示
pd.json_normalize(response["hits"]["hits"])

index_name = "jsquad-neural-search"
query = "日本で梅雨がないのはどこ？"

payload = {
  "query": {
    "multi_match": {
      "query": query,
      "fields": ["question", "context^10"],
      "operator": "and"
    }
  },
  "_source": False,
  "fields": ["question", "answers", "context"],
  "size": 10
}

response = opensearch_client.search(
    index=index_name,
    body=payload
)

pd.json_normalize(response["hits"]["hits"])

index_name = "jsquad-neural-search"
query = "日本で梅雨がないのはどこ？"

payload = {
  "query": {
    "multi_match": {
      "query": query,
      "fields": ["question", "context^10"],
      "operator": "and"
    }
  },
  "ext": {
    "rerank": {
      "query_context": {
         "query_text_path": "query.multi_match.query"
      }
    }
  },
  "_source": False,
  "fields": ["question", "answers", "context"],
  "size": 10,
}

response = opensearch_client.search(
    index=index_name,
    body=payload,
    search_pipeline = rerank_pipeline_id #ベクトル変換を行うパイプラインを指定
)

pd.json_normalize(response["hits"]["hits"])

index_name = "jsquad-neural-search"
query = "日本で梅雨がないのはどこ？"

payload = {
  "query": {
    "hybrid": {
      "queries": [
        {
          "multi_match": {
            "query": query,
            "fields": ["question", "context^10"],
            "operator": "and"
          }        
        },
        {
          "neural": {
            "context_embedding": {
              "query_text": query, # テキストをベクトルに変換し
              "k": 30 # クエリベクトルに近いベクトルのうち上位 30 を返す
            }
          }            
        }
      ]
    }
  },
  "ext": {
    "rerank": {
      "query_context": {
         "query_text": query
      }
    }
  },
  "_source": False,
  "fields": ["question", "answers", "context"],
  "size": 30,
}

response = opensearch_client.search(
    index=index_name,
    body=payload,
    search_pipeline = rerank_pipeline_id #ベクトル変換を行うパイプラインを指定
)

pd.json_normalize(response["hits"]["hits"])