보상 함수 구현

개요

보상 함수(점수 계산기 또는 등급 계산기라고도 함)는 모델 응답을 평가하고 훈련을 위한 피드백 신호를 제공하는 핵심 구성 요소입니다. 이는 모델 응답을 수락하고 보상 점수를 반환하는 Lambda 함수로 구현해야 합니다.

인터페이스 형식

보상 함수는 다음 형식으로 데이터를 수락하고 반환해야 합니다.

훈련을 위한 샘플 입력 샘플

{  
    "messages": [  
        {  
            "role": "user",  
            "content": "Do you have a dedicated security team?"  
        }  
    ],              
   "reference_answer": {  
       "compliant": "No",  
       "explanation": "As an AI developed by Company, I do not have a traditional security team..."  
    }  
}

보상 Lambda용 샘플 페이로드

컨테이너는 다음을 통해 데이터를 Lambda 함수로 전송하기 전에 데이터를 자동으로 변환합니다.

Lambda 함수는 이러한 변환된 형식으로 데이터를 수신합니다.

{    
   "id": "123",  
    "messages": [  
        {  
            "role": "user",  
            "content": "Do you have a dedicated security team?"  
        },  
        {  
            "role": "assistant",  
            "content": "As an AI developed by Amazon, I don not have a dedicated security team..."  
        }  
    ],              
    # Following section will be same as your training dataset sample  
    "reference_answer": {  
        "compliant": "No",  
        "explanation": "As an AI developed by Company, I do not have a traditional security team..."  
    }  
}

보상 Lambda 계약

def lambda_handler(event, context):  
   return lambda_grader(event)  
  
def lambda_grader(samples: list[dict]) -> list[dict]:  
    """  
    Args:  
        samples: List of dictionaries in OpenAI format  
          
        Example input:  
        {     
            "id": "123",  
            "messages": [  
                {  
                    "role": "user",  
                    "content": "Do you have a dedicated security team?"  
                },  
                {  
                    "role": "assistant",  
                    "content": "As an AI developed by Company, I don nott have a dedicated security team..."  
                }  
            ],              
            # This section will be same as your training dataset  
            "reference_answer": {  
                "compliant": "No",  
                "explanation": "As an AI developed by Company, I do not have a traditional security team..."  
            }  
        }  
      
    Returns:  
        List of dictionaries with reward scores:  
        {  
            "id": str,                              # Same id as input sample  
            "aggregate_reward_score": float,        # Overall score for the sample  
            "metrics_list": [                       # OPTIONAL: Component scores  
                {  
                    "name": str,                    # Name of the component score  
                    "value": float,                 # Value of the component score  
                    "type": str                     # "Reward" or "Metric"  
                }  
            ]  
        }  
    """

입력 및 출력 필드

입력 필드

출력 필드

기술적 제약 조건

모범 사례

사용자 지정 보상 함수 사용

태스크별 평가 기준이 있는 경우 사용자 지정 보상 함수를 구현합니다.

IAM 권한

필수 권한

SageMaker 실행 역할에는 Lambda 함수를 간접 호출할 권한이 있어야 합니다. SageMaker 실행 역할에 이 정책을 추가합니다.

{  
  "Version": "2012-10-17",  
  "Statement": [  
    {  
      "Effect": "Allow",  
      "Action": [  
        "lambda:InvokeFunction"  
      ],  
      "Resource": "arn:aws:lambda:region:account-id:function:function-name"  
    }  
  ]  
}

Lambda 실행 역할

Lambda 함수의 실행 역할에는 기본 Lambda 실행 권한이 필요합니다.

{  
  "Version": "2012-10-17",  
  "Statement": [  
    {  
      "Effect": "Allow",  
      "Action": [  
        "logs:CreateLogGroup",  
        "logs:CreateLogStream",  
        "logs:PutLogEvents"  
      ],  
      "Resource": "arn:aws:logs:*:*:*"  
    }  
  ]  
}

추가 권한: Lambda 함수가 다른 AWS 서비스(예: 참조 데이터의 경우 S3, 로깅의 경우 DynamoDB)에 액세스하는 경우 해당 권한을 Lambda 실행 역할에 추가합니다.

예: 평가형 LLM 보상 함수

이 예제에서는 Amazon Bedrock 모델을 평가자로 사용하여 모델 응답을 참조 답변과 비교하여 평가하는 방법을 보여줍니다. 이 Lambda 템플릿은 고객이 평가자의 평가를 처리하기 위한 추론 요청에 대해 Amazon Bedrock에 대한 직접 호출을 구현할 수 있는 프레임워크를 제공합니다. Lambda 함수는 다른 보상 함수와 동일한 입력/출력 계약을 유지 관리합니다.

구현

이 Lambda 함수는 2단계 평가 프로세스를 구현합니다. 즉, lambda_handler는 수신 샘플에서 모델 응답과 참조 응답을 추출한 다음, lambda_graded 함수에서 Amazon Bedrock을 직접 호출하여 이들 간의 시맨틱 유사성을 평가합니다. 이 구현에서는 일시적인 장애에 대한 자동 재시도를 통한 강력한 오류 처리를 포함하며 유연한 참조 응답 형식(문자열 및 구조화된 사전 형식 모두)을 지원합니다.

구현 세부 정보:

"""  
LLM Judge Lambda POC - Working implementation using Amazon Bedrock  
"""  
  
import json  
import time  
import boto3  
  
bedrock_runtime = boto3.client('bedrock-runtime', region_name='us-east-1')  
JUDGE_MODEL_ID = "anthropic.claude-3-5-sonnet-20240620-v1:0"  
SYSTEM_PROMPT = "You must output ONLY a number between 0.0 and 1.0. No explanations, no text, just the number."  
  
JUDGE_PROMPT_TEMPLATE = """Compare the following two responses and rate how similar they are on a scale of 0.0 to 1.0, where:  
- 1.0 means the responses are semantically equivalent (same meaning, even if worded differently)  
- 0.5 means the responses are partially similar  
- 0.0 means the responses are completely different or contradictory  
  
Response A: {response_a}  
  
Response B: {response_b}  
  
Output ONLY a number between 0.0 and 1.0. No explanations."""  
  
  
def lambda_graded(response_a: str, response_b: str, max_retries: int = 3) -> float:  
    """Call Bedrock to compare responses and return similarity score."""  
    prompt = JUDGE_PROMPT_TEMPLATE.format(response_a=response_a, response_b=response_b)  
      
    for attempt in range(max_retries):  
        try:  
            response = bedrock_runtime.converse(  
                modelId=JUDGE_MODEL_ID,  
                messages=[{"role": "user", "content": [{"text": prompt}]}],  
                system=[{"text": SYSTEM_PROMPT}],  
                inferenceConfig={"temperature": 0.0, "maxTokens": 10}  
            )  
            print(f"Bedrock call successful: {response}")  
            output = response['output']['message']['content'][0]['text'].strip()  
            score = float(output)  
            print(f"Score parsed: {score}")  
            return max(0.0, min(1.0, score))  
                  
        except Exception as e:  
            if "ThrottlingException" in str(e) and attempt < max_retries - 1:  
                time.sleep(2 ** attempt)  
            else:  
                print(f"Bedrock call failed: {e}")  
                return None  
    return None  
  
  
def lambda_handler(event, context):  
    """AWS Lambda handler - processes samples from RFTEvalInvoker."""  
    try:  
        samples = event if isinstance(event, list) else [event]  
        results = []  
          
        for sample in samples:  
            sample_id = sample.get("id", "unknown")  
            messages = sample.get("messages", [])  
              
            # Extract assistant response (response A)  
            response_a = ""  
            for msg in messages:  
                if msg.get("role") in ["assistant", "nova_assistant"]:  
                    response_a = msg.get("content", "")  
                    break  
              
            # Extract reference answer from root level (no longer in metadata)  
            reference_answer = sample.get("reference_answer", "")  
              
            # Handle both string and dict reference_answer formats  
            if isinstance(reference_answer, dict):  
                # If reference_answer is a dict, extract the explanation or compliant field  
                response_b = reference_answer.get("explanation", reference_answer.get("compliant", ""))  
            else:  
                response_b = reference_answer  
              
            if not response_a or not response_b:  
                results.append({  
                    "id": sample_id,  
                    "aggregate_reward_score": 0.0,  
                    "metrics_list": [{"name": "similarity_score", "value": 0.0, "type": "Metric"}]  
                })  
                continue  
              
            # Get similarity score  
            score = lambda_graded(response_a, response_b)  
              
            results.append({  
                "id": sample_id,  
                "aggregate_reward_score": score,  
                "metrics_list": [  
                    {  
                        "name": "similarity_score",  
                        "value": score,  
                        "type": "Metric"  
                    }  
                ]  
            })  
          
        return {"statusCode": 200, "body": json.dumps(results)}  
          
    except Exception as e:  
        print(f"Error: {e}")  
        return {"statusCode": 500, "body": json.dumps({"error": str(e)})}

입력 형식

Lambda는 다른 보상 함수와 동일한 입력 형식을 수신합니다.

{  
    "id": "sample-001",  
    "messages": [  
        {  
            "role": "user",  
            "content": "Do you have a dedicated security team?"  
        },  
        {  
            "role": "assistant",  
            "content": "As an AI developed by Amazon, I don't have a dedicated security team..."  
        }  
    ],  
    "reference_answer": {  
        "compliant": "No",  
        "explanation": "As an AI developed by Company, I do not have a traditional security team..."  
    },  
    "my_custom_field": "custom_value"  
}

출력 형식

{  
    "id": "sample-001",  
    "aggregate_reward_score": 0.85,  
    "metrics_list": [  
        {  
            "name": "similarity_score",  
            "value": 0.85,  
            "type": "Metric"  
        }  
    ]  
}

배포 고려 사항

선택한 모델의 기능 및 API 형식에 따라 프롬프트 템플릿 및 추론 파라미터를 조정해야 할 수도 있습니다.

Bedrock 처리량 증가

평가 중에 스로틀링이 발생하는 경우 Bedrock 모델 할당량을 늘립니다.

Lambda의 기본 제공 재시도 로직은 간헐적인 스로틀링을 처리하지만 지속적인 대규모 평가에는 적절한 할당량 증가가 필요합니다.

필수 IAM 정책:

{  
    "Version": "2012-10-17",  
    "Statement": [  
        {  
            "Effect": "Allow",  
            "Action": [  
                "bedrock:InvokeModel"  
            ],  
            "Resource": "arn:aws:bedrock:*::foundation-model/*"  
        }  
    ]  
}