Function evaluate

src/MaxText/rl/evaluate_rl.py:150–233

Computes accuracy and percentage of outputs matching the format. Args: tmvp_config: Configuration object dataset: The evaluation dataset rl_cluster: Model cluster for generation. num_passes: Number of generation passes corr_lst: If True, only include correct respo

(
    tmvp_config,
    dataset,
    rl_cluster,
    num_passes=1,
    corr_lst=False,
    make_lst=False,
)

def evaluate(
    tmvp_config,
    dataset,
    rl_cluster,
    num_passes=1,
    corr_lst=False,
    make_lst=False,
):
  """
  Computes accuracy and percentage of outputs matching the format.

  Args:
      tmvp_config: Configuration object
      dataset: The evaluation dataset
      rl_cluster: Model cluster for generation.
      num_passes: Number of generation passes
      corr_lst: If True, only include correct responses in the list
      make_lst: If True, return a list of (question, answer, responses)

  Returns:
      Tuple of statistics and optionally the response list
  """
  response_lst = []
  corr = 0
  partially_corr = 0
  corr_format = 0
  total = 0

  for batch in tqdm(dataset):
    answers = batch["answer"]
    questions = batch["question"]
    prompts = batch["prompts"]

    # Generate responses for all prompts in the batch
    multiple_call_responses = generate_responses(
        tmvp_config=tmvp_config,
        prompts=prompts,
        rl_cluster=rl_cluster,
        num_passes=num_passes,
    )

    # Score each question-answer pair
    for question, responses, answer in zip(questions, multiple_call_responses, answers):
      is_correct, is_partially_correct, has_correct_format = score_responses(
          tmvp_config=tmvp_config,
          question=question,
          responses=responses,
          answer=answer,
      )

      # Update counters
      if is_correct:
        corr += 1
        if corr_lst and make_lst:
          response_lst.append((question, answer, responses))
      else:
        if not corr_lst and make_lst:
          response_lst.append((question, answer, responses))