ベクトル検索の実装 (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 awswrangler as wr
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
    embedding_model_id_on_hf = "BAAI/bge-m3"
    embedding_model_name = embedding_model_id_on_hf.lower().replace("/", "-")
    realtime_inference_endpoint_type = "instance"

    embedding_endpoint_name_prefix = f"{embedding_model_name}-{realtime_inference_endpoint_type}"
    response = search_sagemaker_inference_endpoint(embedding_endpoint_name_prefix, sagemaker_region)

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

        hub = {
          "HF_MODEL_ID": embedding_model_id_on_hf,
          "HF_TASK":"feature-extraction", # NLP task you want to use for predictions
        }
        
        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.")
        embedding_model = HuggingFaceModel(
            name=sagemaker.utils.name_from_base(embedding_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 {embedding_model_id_on_hf} on a realtime {realtime_inference_endpoint_type} endpoint.")
        embedding_model.deploy(
            endpoint_name=sagemaker.utils.name_from_base(embedding_endpoint_name_prefix),
            initial_instance_count=1,
            instance_type=instance_type
        )

        embedding_endpoint_name = embedding_model.endpoint_name

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

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

%%time

payload = {
    "inputs": ["hello world!"]
}
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=embedding_endpoint_name,
    ContentType="application/json",
    Accept="application/json",
    Body=body
)

embeddings = eval(response['Body'].read().decode('utf-8'))

print("embedding:")
print(np.array(embeddings[0]))
print("dimension:")
print(np.shape(embeddings[0]))

%%time
dataset_dir = "./dataset/jsquad/"
%mkdir -p $dataset_dir
!curl -L -s -o $dataset_dir/valid.json https://github.com/yahoojapan/JGLUE/raw/main/datasets/jsquad-v1.3/valid-v1.3.json 
#!curl -L -s -o $dataset_dir/train.json https://github.com/yahoojapan/JGLUE/raw/main/datasets/jsquad-v1.3/train-v1.3.json 

%%time
import pandas as pd
import json

def squad_json_to_dataframe(input_file_path, record_path=["data", "paragraphs", "qas", "answers"]):
    file = json.loads(open(input_file_path).read())
    m = pd.json_normalize(file, record_path[:-1])
    r = pd.json_normalize(file, record_path[:-2])

    idx = np.repeat(r["context"].values, r.qas.str.len())
    m["context"] = idx
    m["answers"] = m["answers"]
    m["answers"] = m["answers"].apply(lambda x: np.unique(pd.json_normalize(x)["text"].to_list()))
    return m[["id", "question", "context", "answers"]]

valid_filename = f"{dataset_dir}/valid.json"
valid_df = squad_json_to_dataframe(valid_filename)

#train_filename = f"{dataset_dir}/train.json"
#train_df = squad_json_to_dataframe(train_filename)

valid_df

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()

payload = {
  "mappings": {
    "properties": {
      "id": {"type": "keyword"},
      "question": {"type": "text", "analyzer": "custom_sudachi_analyzer"},
      "context":  {"type": "text", "analyzer": "custom_sudachi_analyzer"},
      "answers":  {"type": "text", "analyzer": "custom_sudachi_analyzer"},
      "question_embedding": {
        "type": "knn_vector",
        "dimension": 1024,
        "space_type": "l2",
        "method": {
          "name": "hnsw",
          "engine": "faiss",
        }
      },
      "context_embedding": {
        "type": "knn_vector",
        "dimension": 1024,
        "space_type": "l2",
        "method": {
          "name": "hnsw",
          "engine": "faiss",
        },
      }
    }
  },
  "settings": {
    "index.knn": True,
    "index.number_of_shards": 1,
    "index.number_of_replicas": 0,
    "analysis": {
      "analyzer": {
        "custom_sudachi_analyzer": {
          "char_filter": ["icu_normalizer"],
          "filter": [
              "sudachi_normalizedform",
              "custom_sudachi_part_of_speech"
          ],
          "tokenizer": "sudachi_tokenizer",
          "type": "custom"
        }
      },
      "filter": {
        "custom_sudachi_part_of_speech": {
          "type": "sudachi_part_of_speech",
          "stoptags": ["感動詞,フィラー","接頭辞","代名詞","副詞","助詞","助動詞","動詞,一般,*,*,*,終止形-一般","名詞,普通名詞,副詞可能"]
        }
      }
    }
  }
}
# インデックス名を指定
index_name = "jsquad-knn"

try:
    # 既に同名のインデックスが存在する場合、いったん削除を行う
    print("# delete index")
    response = opensearch_client.indices.delete(index=index_name)
    print(json.dumps(response, indent=2))
except Exception as e:
    print(e)

# インデックスを作成
response = opensearch_client.indices.create(index_name, body=payload)
response

%%time

def get_df_with_embeddings(input_df, field_mappings, embedding_endpoint_name, sagemaker_region, batch_size):
    output_df = pd.DataFrame([]) #create empty dataframe
    df_list = np.array_split(input_df, input_df.shape[0]/batch_size)
    for df in tqdm(df_list):
        index = df.index #backup index number
        df_with_embeddings = df
        for field_mapping in field_mappings:
            input_field_name = field_mapping["InputFieldName"]
            embedding_field_name = field_mapping["EmbeddingFieldName"]
            #index = df.index #backup index number
            payload = {
                "inputs": df_with_embeddings[input_field_name].values.tolist()
            }
            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=embedding_endpoint_name,
                ContentType="application/json",
                Accept="application/json",
                Body=body
            )
            embeddings = eval(response['Body'].read().decode('utf-8'))
            df_with_embeddings = pd.concat([df_with_embeddings.reset_index(drop=True), pd.Series(embeddings,name=embedding_field_name).reset_index(drop=True)],axis=1) #join embedding results to source dataframe
            df_with_embeddings = df_with_embeddings.set_index(index) #restore index number
        output_df = pd.concat([output_df, df_with_embeddings])
    return output_df

valid_df_with_embeddings = get_df_with_embeddings(
    input_df=valid_df,
    field_mappings=[
        {"InputFieldName": "question", "EmbeddingFieldName": "question_embedding"},
        {"InputFieldName": "context", "EmbeddingFieldName": "context_embedding"},
    ],
    embedding_endpoint_name=embedding_endpoint_name,
    sagemaker_region=sagemaker_region,
    batch_size=20
)

valid_df_with_embeddings

%%time
index_name = "jsquad-knn"
response = wr.opensearch.index_df(
    client=opensearch_client,
    df=valid_df_with_embeddings,
    use_threads=True,
    id_keys=["id"],
    index=index_name,
    bulk_size=200, # 200 件ずつ書き込み
    refresh=False,
)

response["success"] == valid_df["id"].count()

index_name = "jsquad-knn"
query = "日本 梅雨 ない どこ"

payload = {
  "query": {
    "match": {
      "question": {
        "query": query,
        "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-knn"
query = "日本で梅雨がない場所は？"

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

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

response

index_name = "jsquad-knn"
query = "日本で梅雨がない場所は？"

payload = {
  "query": {
    "match": {
      "question": {
        "query": query,
        "operator": "or",
        "minimum_should_match": "75%"
      }
    }
  },
  "_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-knn"

query = "日本で梅雨がない場所は？"

def text_to_embedding(text, region_name, embedding_endpoint_name):
    payload = {
        "inputs": [
            query
        ]
    }
    body = bytes(json.dumps(payload), 'utf-8')
    sagemaker_runtime_client = boto3.client("sagemaker-runtime", region_name=region_name)
    response = sagemaker_runtime_client.invoke_endpoint(
        EndpointName=embedding_endpoint_name,
        ContentType="application/json",
        Accept="application/json",
        Body=body
    )
    embeddings = eval(response['Body'].read().decode('utf-8'))
    return embeddings[0]

vector = text_to_embedding(text=query, region_name=sagemaker_region, embedding_endpoint_name=embedding_endpoint_name)
k = 10

payload = {
  "query": {
    "knn": {
      "question_embedding": {
        "vector": vector,
        "k": k
      }
    }
  },
  "_source": False,
  "fields": ["question", "answers", "context"],
  "size": k
}

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

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

k = 10

payload = {
  "query": {
    "knn": {
      "question_embedding": {
        "vector": vector,
        "min_score": 0.7
      }
    }
  },
  "_source": False,
  "fields": ["question", "answers", "context"],
  "size": k,
  
}

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

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

k = 10

payload = {
  "query": {
    "knn": {
      "question_embedding": {
        "vector": vector,
        "max_distance": 0.3
      }
    }
  },
  "_source": False,
  "fields": ["question", "answers", "context"],
  "size": k,
  
}

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

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