KBR's Stacked Thoughts

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
import functools

# 최대한 간결한 ver.
terse_guidance = "Answer the following question in a single sentence in Korean, or as close to that as possible."
#  필요하다면 출처를 인용하지만, 답변에 꼭 필요한 만큼만 사용하는 간단한 ver.
moderate_guidance = "Provide a brief answer to the following question in Korean, use a citation if necessary, but only enough to answer the question."
# 문서를 인용하고 가능한 한 배경 정보도 함께 제공하는 구체적인 ver.
cited_guidance = "Provide a thorough, detailed answer to the following question in Korean, citing the document and supplying additional background information as much as possible."
guidance_options = {
    'Terse': terse_guidance,
    'Moderate': moderate_guidance,
    'Cited': cited_guidance,
}

questions = [
    # 더 많은 질문을 평가할수록 시간이 더 소요되지만, 신뢰도 높은 결과를 얻을 수 있다. 
    # 실제 production 시스템에서는 수백 개의 질문으로 복잡한 시스템을 평가하기도 한다. 
    "What metric(s) are used to evaluate long context performance?",
    "How does the model perform on code tasks?",
    "How many layers does it have?",
    "Why is it called Gemini?",
]

# 평가할 질문이 없을 경우
if not questions:
  raise NotImplementedError('Add some questions to evaluate!')

@functools.cache
def answer_question(question: str, guidance: str = '') -> str:
  """업로드된 문서와 instruction을 기반으로 질문에 대한 답변 생성하기"""
  config = types.GenerateContentConfig(
      temperature=0.0,
      system_instruction=guidance,
  )
  response = client.models.generate_content(
      model='gemini-2.0-flash',
      config=config,
      contents=[question, document_file],
  )
  
  return response.text

1
2
3
# 첫 번째 질문-간결하게: What metric(s) are used to evaluate long context performance?
answer1 = answer_question(questions[0], terse_guidance)
Markdown(answer1)

1
2
3
# 두 번째 질문-간단하게: "How does the model perform on code tasks?"
answer2 = answer_question(questions[1], moderate_guidance)
Markdown(answer2)

1
2
3
# 세 번째 질문-자세하게: "How many layers does it have?"
answer3 = answer_question(questions[2], cited_guidance)
Markdown(answer3)

1
2
3
# 네 번째 질문-자세하게: "Why is it called Gemini?"
answer4 = answer_question(questions[3], cited_guidance)
Markdown(answer4)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
import enum

QA_PROMPT = """\
 (생략)

# User Inputs and AI-generated Response
## User Inputs
### Prompt
{prompt}

## AI-generated Response
{response}
"""

# 답변 평가 점수를 나타내는 enum 클래스
class AnswerRating(enum.Enum):
  VERY_GOOD = '5'
  GOOD = '4'
  OK = '3'
  BAD = '2'
  VERY_BAD = '1'


@functools.cache
def eval_answer(prompt, ai_response, n=1):
  """주어진 질문과 AI 응답 평가하기"""
  chat = client.chats.create(model='gemini-2.0-flash')

  # 평가용 전체 텍스트 생성 요청
  response = chat.send_message(
      message=QA_PROMPT.format(prompt=[prompt, document_file], response=ai_response)
  )
  verbose_eval = response.text


  # 결과를 구조화된 enum 형태로 변환하기 위한 설정
  structured_output_config = types.GenerateContentConfig(
      response_mime_type="text/x.enum",
      response_schema=AnswerRating,
  )
  # 최종 점수만 추출
  response = chat.send_message(
      message="Convert the final score.",
      config=structured_output_config,
  )
  structured_eval = response.parsed

  return verbose_eval, structured_eval

# 질문에 대해 답변을 평가하고 결과 출력
text_eval, struct_eval = eval_answer(prompt=questions[0], ai_response=answer1)
display(Markdown(text_eval))
print(struct_eval)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
import collections
import itertools

# 각 작업을 반복할 횟수 설정
# 반복 횟수를 늘리면 시간이 오래 걸리지만, 오차를 줄이고 평균 결과를 얻을 수 있다 
NUM_ITERATIONS = 2

# 점수를 저장할 딕셔너리
scores = collections.defaultdict(int)

# 생성된 응답과 평가를 저장할 딕셔너리
responses = collections.defaultdict(list)

for question in questions:
  display(Markdown(f'## {question}'))
  for guidance, guide_prompt in guidance_options.items():

    for n in range(NUM_ITERATIONS):
      # 답변 생성
      answer = answer_question(question, guide_prompt)

      # 답변 평가 (prompt와 instruction은 전달하지 않음 주의!) 
      written_eval, struct_eval = eval_answer(question, answer, n)
      print(f'{guidance}: {struct_eval}')

      # 평가 점수 저장
      scores[guidance] += int(struct_eval.value)

      # 생성된 응답과 평가 결과를 저장
      responses[(guidance, question)].append((answer, written_eval))

1
2
3
4
5
6
7
for guidance, score in scores.items():
  # 평균 점수 계산: 총 점수 / (반복 횟수 * 질문 수)
  avg_score = score / (NUM_ITERATIONS * len(questions))

  # 평균 점수를 가장 가까운 정수로 반올림한 뒤, AnswerRating enum으로 변환
  nearest = AnswerRating(str(round(avg_score)))
  print(f'{guidance}: {avg_score:.2f} - {nearest.name}')

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
QA_PAIRWISE_PROMPT = """ (생략)
# User Inputs and AI-generated Responses
## User Inputs
### Prompt
{prompt}

# AI-generated Response

### Response A
{baseline_model_response}

### Response B
{response}
"""

# 평가 결과로 반환될 값 (A가 더 낫다 / B가 더 낫다 / 동일하다)
class AnswerComparison(enum.Enum):
  A = 'A'
  SAME = 'SAME'
  B = 'B'

@functools.cache
def eval_pairwise(prompt, response_a, response_b, n=1):
  """같은 질문에 대한 두 개의 답변 중 어떤 것이 더 나은지 평가"""

  # Gemini 모델 채팅 세션 생성
  chat = client.chats.create(model='gemini-2.0-flash')

  # 전체 평가 텍스트 생성 요청
  response = chat.send_message(
      message=QA_PAIRWISE_PROMPT.format(
          prompt=[prompt, document_file],
          baseline_model_response=response_a,
          response=response_b)
  )
  verbose_eval = response.text

  # 구조화된 출력 형식으로 점수 변환
  structured_output_config = types.GenerateContentConfig(
      response_mime_type="text/x.enum",
      response_schema=AnswerComparison,
  )
  # 최종 평가 결과를 enum 형태로 요청
  response = chat.send_message(
      message="Convert the final score.",
      config=structured_output_config,
  )
  structured_eval = response.parsed

  return verbose_eval, structured_eval

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
# 실제 실행: 하나의 질문에 대해 두 가지 가이드라인으로 생성된 답변 비교
question = questions[0]
answer_a = answer_question(question, terse_guidance)  # 간결한 답변
answer_b = answer_question(question, cited_guidance)  # 인용 포함 상세 답변

# pairwise 평가 실행
text_eval, struct_eval = eval_pairwise(
    prompt=question,
    response_a=answer_a,
    response_b=answer_b,
)

# 평가에 대한 설명과 최종 선택 결과 출력
display(Markdown(text_eval))
print(struct_eval)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
@functools.total_ordering
class QAGuidancePrompt:
  """질문-응답에서 사용하는 prompt 또는 instruction"""
    
  def __init__(self, prompt, questions, n_comparisons=NUM_ITERATIONS):
    """prompt 인스턴스를 생성
       평가에 사용할 질문 리스트와 각 질문당 평가를 몇 번 수행할지 결정"""
    self.prompt = prompt
    self.questions = questions
    self.n = n_comparisons

  def __str__(self):
    """문자열로 출력할 때 프롬프트 내용이 보이도록 설정"""
    return self.prompt

  def _compare_all(self, other):
    """모든 질문에 대해 n회씩 비교 수행하여 평균적인 비교 결과를 반환한다.
       이 값이 > 0이면 self가 더 낫고, < 0이면 other가 낫다"""
    results = [self._compare_n(other, q) for q in questions]
    mean = sum(results) / len(results)
    return round(mean)

  def _compare_n(self, other, question):
    """하나의 질문에 대해 n회 평가를 수행하고 평균 결과 반"""
    results = [self._compare(other, question, n) for n in range(self.n)]
    mean = sum(results) / len(results)
    return mean

  def _compare(self, other, question, n=1):
    """한 번의 질문 평가에서 두 프롬프트 중 무엇이 더 나은지 비교
       Gemini로부터 응답을 생성하고 pairwise 평가 수행"""
    answer_a = answer_question(question, self.prompt)
    answer_b = answer_question(question, other.prompt)

    _, result = eval_pairwise(
        prompt=question,
        response_a=answer_a,
        response_b=answer_b,
        n=n,  # Cache buster (캐시 무효화를 위한 인자)
    )
    # print(f'q[{question}], a[{self.prompt[:20]}...], b[{other.prompt[:20]}...]: {result}')

    # 평가 결과 enum을 정수로 변환 (A가 낫다 = 1, B가 낫다 = -1, 동일 = 0)
    if result is AnswerComparison.A:
      return 1
    elif result is AnswerComparison.B:
      return -1
    else:
      return 0

  def __eq__(self, other):
    """두 프롬프트가 동등한지 판단 (pairwise 평가 평균이 0인지로 판단)"""
    if not isinstance(other, QAGuidancePrompt):
      return NotImplemented

    return self._compare_all(other) == 0

  def __lt__(self, other):
    """self가 other보다 열등한지 판단 (평균 비교 결과가 음수이면 True)"""
    if not isinstance(other, QAGuidancePrompt):
      return NotImplemented

    return self._compare_all(other) < 0

1
2
3
4
5
6
7
8
9
10
11
12
# 각 프롬프트를 평가 가능한 객체로 설정
terse_prompt = QAGuidancePrompt(terse_guidance, questions)
moderate_prompt = QAGuidancePrompt(moderate_guidance, questions)
cited_prompt = QAGuidancePrompt(cited_guidance, questions)

# 프롬프트 정렬 (reverse = True로 설정해서 가장 우수한 프롬프트가 첫 번째로 오도록)
sorted_results = sorted([terse_prompt, moderate_prompt, cited_prompt], reverse=True)
for i, p in enumerate(sorted_results):
  if i:
    print('---')

  print(f'#{i+1}: {p}')