Reranking search results using Cohere Rerank on Amazon Bedrock
출처: 문서
본문
Amazon Bedrock에서 Cohere Rerank로 검색 결과 리랭킹하기
이 튜토리얼에서는 Amazon Bedrock에서 호스팅되는 Cohere Rerank 모델을 사용해 Amazon OpenSearch Service와 자체 관리형 OpenSearch에서 검색 결과 리랭킹을 구현하는 방법을 알아요.
리랭킹 파이프라인은 검색 결과를 리랭킹하면서, 검색 쿼리를 기준으로 검색 결과 속 각 문서에 관련성 점수를 부여해요. 관련성 점수는 크로스 인코더 모델이 계산해요.
your_ 접두사로 시작하는 플레이스홀더는 여러분의 값으로 바꾸세요.
{: .note}
사전 준비: Amazon Bedrock에서 모델 테스트하기
모델을 사용하기 전에 다음 코드로 Amazon Bedrock에서 테스트해 보세요:
import json
import boto3
bedrock_region = "your_bedrock_model_region_like_us-west-2"
bedrock_runtime_client = boto3.client("bedrock-runtime", region_name=bedrock_region)
modelId = "cohere.rerank-v3-5:0"
contentType = "application/json"
accept = "*/*"
body = json.dumps({
"query": "What is the capital city of America?",
"documents": [
"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."
],
"api_version": 2
})
response = bedrock_runtime_client.invoke_model(
modelId=modelId,
contentType=contentType,
accept=accept,
body=body
)
results = json.loads(response.get('body').read())["results"]
print(json.dumps(results, indent=2))
{% include copy.html %}
응답에는 관련성 점수 순으로 정렬된 리랭킹 결과가 담겨 있어요:
[
{
"index": 2,
"relevance_score": 0.7190094
},
{
"index": 0,
"relevance_score": 0.32418242
},
{
"index": 1,
"relevance_score": 0.07456104
},
{
"index": 3,
"relevance_score": 0.06124987
}
]
결과를 인덱스 순으로 정렬하려면 다음 코드를 사용하세요:
print(json.dumps(sorted(results, key=lambda x: x['index']), indent=2))
{% include copy.html %}
정렬된 결과는 다음과 같아요:
[
{
"index": 0,
"relevance_score": 0.32418242
},
{
"index": 1,
"relevance_score": 0.07456104
},
{
"index": 2,
"relevance_score": 0.7190094
},
{
"index": 3,
"relevance_score": 0.06124987
}
]
1단계: 커넥터 생성 및 모델 등록하기
모델용 커넥터를 만들려면 다음 요청을 보내세요.
자체 관리형 OpenSearch를 사용한다면 AWS 자격 증명을 입력하세요:
POST /_plugins/_ml/connectors/_create
{
"name": "Amazon Bedrock Cohere rerank model",
"description": "Test connector for Amazon Bedrock Cohere rerank model",
"version": 1,
"protocol": "aws_sigv4",
"credential": {
"access_key": "your_access_key",
"secret_key": "your_secret_key",
"session_token": "your_session_token"
},
"parameters": {
"service_name": "bedrock",
"endpoint": "bedrock-runtime",
"region": "your_bedrock_model_region_like_us-west-2",
"model_name": "cohere.rerank-v3-5:0",
"api_version": 2
},
"actions": [
{
"action_type": "PREDICT",
"method": "POST",
"url": "https://${parameters. endpoint}.${parameters.region}.amazonaws.com/model/${parameters.model_name}/invoke",
"headers": {
"x-amz-content-sha256": "required",
"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 + '", "documents": ' + textDocsBuilder.toString() + ' }';
return '{"parameters": ' + parameters + '}';
""",
"request_body": """
{
"documents": ${parameters.documents},
"query": "${parameters.query}",
"api_version": ${parameters.api_version}
}
""",
"post_process_function": """
if (params.results == null || params.results.length == 0) {
throw new IllegalArgumentException("Post process function input is empty.");
}
def outputs = params.results;
def relevance_scores = new Double[outputs.length];
for (int i=0; i<outputs.length; i++) {
def index = new BigDecimal(outputs[i].index.toString()).intValue();
relevance_scores[index] = outputs[i].relevance_score;
}
def resultBuilder = new StringBuilder('[');
for (int i=0; i<relevance_scores.length; i++) {
resultBuilder.append(' {"name": "similarity", "data_type": "FLOAT32", "shape": [1],');
resultBuilder.append('"data": [');
resultBuilder.append(relevance_scores[i]);
resultBuilder.append(']}');
if (i<outputs.length - 1) {
resultBuilder.append(',');
}
}
resultBuilder.append(']');
return resultBuilder.toString();
"""
}
]
}
{% include copy-curl.html %}
Amazon OpenSearch Service를 사용한다면 Amazon Bedrock에 대한 액세스를 허용하는 AWS Identity and Access Management(IAM) 역할의 Amazon Resource Name(ARN)을 제공할 수 있어요:
POST /_plugins/_ml/connectors/_create
{
"name": "Amazon Bedrock Cohere rerank model",
"description": "Test connector for Amazon Bedrock Cohere rerank model",
"version": 1,
"protocol": "aws_sigv4",
"credential": {
"roleArn": "your_role_arn_which_allows_access_to_bedrock_model"
},
"parameters": {
"service_name": "bedrock",
"endpoint": "bedrock-runtime",
"region": "your_bedrock_model_region_like_us-west-2",
"model_name": "cohere.rerank-v3-5:0",
"api_version": 2
},
"actions": [
{
"action_type": "PREDICT",
"method": "POST",
"url": "https://${parameters. endpoint}.${parameters.region}.amazonaws.com/model/${parameters.model_name}/invoke",
"headers": {
"x-amz-content-sha256": "required",
"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 + '", "documents": ' + textDocsBuilder.toString() + ' }';
return '{"parameters": ' + parameters + '}';
""",
"request_body": """
{
"documents": ${parameters.documents},
"query": "${parameters.query}",
"api_version": ${parameters.api_version}
}
""",
"post_process_function": """
if (params.results == null || params.results.length == 0) {
throw new IllegalArgumentException("Post process function input is empty.");
}
def outputs = params.results;
def relevance_scores = new Double[outputs.length];
for (int i=0; i<outputs.length; i++) {
def index = new BigDecimal(outputs[i].index.toString()).intValue();
relevance_scores[index] = outputs[i].relevance_score;
}
def resultBuilder = new StringBuilder('[');
for (int i=0; i<relevance_scores.length; i++) {
resultBuilder.append(' {"name": "similarity", "data_type": "FLOAT32", "shape": [1],');
resultBuilder.append('"data": [');
resultBuilder.append(relevance_scores[i]);
resultBuilder.append(']}');
if (i<outputs.length - 1) {
resultBuilder.append(',');
}
}
resultBuilder.append(']');
return resultBuilder.toString();
"""
}
]
}
자세한 내용은 AWS 문서를 참고하세요.
응답에서 받은 커넥터 ID로 모델을 등록하고 배포하세요:
POST /_plugins/_ml/models/_register?deploy=true
{
"name": "Amazon Bedrock Cohere rerank model",
"function_name": "remote",
"description": "test rerank model",
"connector_id": "your_connector_id"
}
{% include copy-curl.html %}
응답에 있는 모델 ID를 메모해 두세요. 이후 단계에서 사용하게 돼요.
Predict API를 사용해 모델을 테스트해 보세요:
POST _plugins/_ml/models/your_model_id/_predict
{
"parameters": {
"query": "What is the capital city of America?",
"documents": [
"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."
]
}
}
{% include copy-curl.html %}
또는 다음과 같이 모델을 테스트할 수도 있어요:
POST _plugins/_ml/_predict/text_similarity/your_model_id
{
"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."
]
}
{% include copy-curl.html %}
커넥터의 pre_process_function은 입력을 앞서 보여 준 파라미터에 필요한 형식으로 변환해요.
기본적으로 Amazon Bedrock Rerank API의 출력은 다음과 같은 형식이에요:
[
{
"index": 2,
"relevance_score": 0.7190094
},
{
"index": 0,
"relevance_score": 0.32418242
},
{
"index": 1,
"relevance_score": 0.07456104
},
{
"index": 3,
"relevance_score": 0.06124987
}
]
커넥터의 post_process_function은 모델 출력을 rerank 프로세서가 해석할 수 있는 형식으로 변환하고 결과를 인덱스 순으로 정렬해요. 변환된 형식은 다음과 같아요:
{
"inference_results": [
{
"output": [
{
"name": "similarity",
"data_type": "FLOAT32",
"shape": [
1
],
"data": [
0.32418242
]
},
{
"name": "similarity",
"data_type": "FLOAT32",
"shape": [
1
],
"data": [
0.07456104
]
},
{
"name": "similarity",
"data_type": "FLOAT32",
"shape": [
1
],
"data": [
0.7190094
]
},
{
"name": "similarity",
"data_type": "FLOAT32",
"shape": [
1
],
"data": [
0.06124987
]
}
],
"status_code": 200
}
]
}
응답에는 similarity 객체 네 개가 담겨 있어요. 각 similarity 객체의 data 배열에는 쿼리를 기준으로 한 각 문서의 관련성 점수가 들어 있어요. similarity 객체는 입력 문서 순서대로 제공되는데, 즉 첫 번째 객체가 첫 번째 문서에 해당해요. 이는 문서를 관련성 점수 순으로 정렬하는 Cohere Rerank 모델의 기본 출력과는 다른 부분이에요. 문서 순서는 connector.post_process.cohere.rerank 후처리 함수에서 변경되어, 출력이 리랭킹 파이프라인과 호환되도록 만들어져요.
2단계: 리랭킹 파이프라인 구성하기
다음 단계에 따라 리랭킹 파이프라인을 구성하세요.
2.1단계: 테스트 데이터 수집하기
테스트 데이터를 수집하려면 다음 벌크 요청을 보내세요:
POST _bulk
{ "index": { "_index": "my-test-data" } }
{ "passage_text" : "Carson City is the capital city of the American state of Nevada." }
{ "index": { "_index": "my-test-data" } }
{ "passage_text" : "The Commonwealth of the Northern Mariana Islands is a group of islands in the Pacific Ocean. Its capital is Saipan." }
{ "index": { "_index": "my-test-data" } }
{ "passage_text" : "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." }
{ "index": { "_index": "my-test-data" } }
{ "passage_text" : "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." }
{% include copy-curl.html %}
2.2단계: 리랭킹 파이프라인 생성하기
Cohere Rerank 모델로 리랭킹 파이프라인을 생성하세요:
PUT /_search/pipeline/rerank_pipeline_bedrock
{
"description": "Pipeline for reranking with Bedrock Cohere rerank model",
"response_processors": [
{
"rerank": {
"ml_opensearch": {
"model_id": "your_model_id_created_in_step1"
},
"context": {
"document_fields": ["passage_text"]
}
}
}
]
}
{% include copy-curl.html %}
document_fields에 여러 필드 이름을 지정하면 모든 필드의 값을 먼저 이어 붙인 다음 리랭킹이 수행돼요.
{: .note}
2.3단계: 리랭킹 테스트하기
반환되는 결과 수를 제한하려면 size 파라미터를 지정할 수 있어요. 예를 들어 상위 두 개 문서를 반환하려면 "size": 2로 설정하세요.
먼저 리랭킹 파이프라인을 사용하지 않고 쿼리를 테스트해 보세요:
POST my-test-data/_search
{
"query": {
"match": {
"passage_text": "What is the capital city of America?"
}
},
"highlight": {
"pre_tags": ["<strong>"],
"post_tags": ["</strong>"],
"fields": {"passage_text": {}}
},
"_source": false,
"fields": ["passage_text"]
}
{% include copy-curl.html %}
응답의 첫 번째 문서는 Carson City is the capital city of the American state of Nevada인데, 이는 잘못된 결과예요:
{
"took": 2,
"timed_out": false,
"_shards": {
"total": 1,
"successful": 1,
"skipped": 0,
"failed": 0
},
"hits": {
"total": {
"value": 4,
"relation": "eq"
},
"max_score": 2.5045562,
"hits": [
{
"_index": "my-test-data",
"_id": "1",
"_score": 2.5045562,
"fields": {
"passage_text": [
"Carson City is the capital city of the American state of Nevada."
]
},
"highlight": {
"passage_text": [
"Carson <strong>City</strong> <strong>is</strong> <strong>the</strong> <strong>capital</strong> <strong>city</strong> <strong>of</strong> <strong>the</strong> American state <strong>of</strong> Nevada."
]
}
},
{
"_index": "my-test-data",
"_id": "2",
"_score": 0.5807494,
"fields": {
"passage_text": [
"The Commonwealth of the Northern Mariana Islands is a group of islands in the Pacific Ocean. Its capital is Saipan."
]
},
"highlight": {
"passage_text": [
"<strong>The</strong> Commonwealth <strong>of</strong> <strong>the</strong> Northern Mariana Islands <strong>is</strong> a group <strong>of</strong> islands in <strong>the</strong> Pacific Ocean.",
"Its <strong>capital</strong> <strong>is</strong> Saipan."
]
}
},
{
"_index": "my-test-data",
"_id": "3",
"_score": 0.5261191,
"fields": {
"passage_text": [
"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."
]
},
"highlight": {
"passage_text": [
"(also known as simply Washington or D.C., and officially as <strong>the</strong> District <strong>of</strong> Columbia) <strong>is</strong> <strong>the</strong> <strong>capital</strong>",
"<strong>of</strong> <strong>the</strong> United States.",
"It <strong>is</strong> a federal district."
]
}
},
{
"_index": "my-test-data",
"_id": "4",
"_score": 0.5083029,
"fields": {
"passage_text": [
"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."
]
},
"highlight": {
"passage_text": [
"<strong>Capital</strong> punishment (<strong>the</strong> death penalty) has existed in <strong>the</strong> United States since beforethe United States",
"As <strong>of</strong> 2017, <strong>capital</strong> punishment <strong>is</strong> legal in 30 <strong>of</strong> <strong>the</strong> 50 states."
]
}
}
]
}
}
다음으로 리랭킹 파이프라인을 사용해 쿼리를 테스트해 보세요:
POST my-test-data/_search?search_pipeline=rerank_pipeline_bedrock
{
"query": {
"match": {
"passage_text": "What is the capital city of America?"
}
},
"ext": {
"rerank": {
"query_context": {
"query_text": "What is the capital city of America?"
}
}
},
"highlight": {
"pre_tags": ["<strong>"],
"post_tags": ["</strong>"],
"fields": {"passage_text": {}}
},
"_source": false,
"fields": ["passage_text"]
}
{% include copy-curl.html %}
응답의 첫 번째 문서는 "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."인데, 이번에는 올바른 결과예요:
{
"took": 2,
"timed_out": false,
"_shards": {
"total": 1,
"successful": 1,
"skipped": 0,
"failed": 0
},
"hits": {
"total": {
"value": 4,
"relation": "eq"
},
"max_score": 0.7190094,
"hits": [
{
"_index": "my-test-data",
"_id": "3",
"_score": 0.7190094,
"fields": {
"passage_text": [
"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."
]
},
"highlight": {
"passage_text": [
"(also known as simply Washington or D.C., and officially as <strong>the</strong> District <strong>of</strong> Columbia) <strong>is</strong> <strong>the</strong> <strong>capital</strong>",
"<strong>of</strong> <strong>the</strong> United States.",
"It <strong>is</strong> a federal district."
]
}
},
{
"_index": "my-test-data",
"_id": "1",
"_score": 0.32418242,
"fields": {
"passage_text": [
"Carson City is the capital city of the American state of Nevada."
]
},
"highlight": {
"passage_text": [
"Carson <strong>City</strong> <strong>is</strong> <strong>the</strong> <strong>capital</strong> <strong>city</strong> <strong>of</strong> <strong>the</strong> American state <strong>of</strong> Nevada."
]
}
},
{
"_index": "my-test-data",
"_id": "2",
"_score": 0.07456104,
"fields": {
"passage_text": [
"The Commonwealth of the Northern Mariana Islands is a group of islands in the Pacific Ocean. Its capital is Saipan."
]
},
"highlight": {
"passage_text": [
"<strong>The</strong> Commonwealth <strong>of</strong> <strong>the</strong> Northern Mariana Islands <strong>is</strong> a group <strong>of</strong> islands in <strong>the</strong> Pacific Ocean.",
"Its <strong>capital</strong> <strong>is</strong> Saipan."
]
}
},
{
"_index": "my-test-data",
"_id": "4",
"_score": 0.06124987,
"fields": {
"passage_text": [
"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."
]
},
"highlight": {
"passage_text": [
"<strong>Capital</strong> punishment (<strong>the</strong> death penalty) has existed in <strong>the</strong> United States since beforethe United States",
"As <strong>of</strong> 2017, <strong>capital</strong> punishment <strong>is</strong> legal in 30 <strong>of</strong> <strong>the</strong> 50 states."
]
}
}
]
},
"profile": {
"shards": []
}
}
쿼리를 두 번 작성하지 않으려면 query_text 대신 query_text_path를 다음과 같이 사용하세요:
POST my-test-data/_search?search_pipeline=rerank_pipeline_bedrock
{
"query": {
"match": {
"passage_text": "What is the capital city of America?"
}
},
"ext": {
"rerank": {
"query_context": {
"query_text_path": "query.match.passage_text.query"
}
}
},
"highlight": {
"pre_tags": ["<strong>"],
"post_tags": ["</strong>"],
"fields": {"passage_text": {}}
},
"_source": false,
"fields": ["passage_text"]
}
{% include copy-curl.html %}