Function evaluate_model

src/MaxText/examples/sft_train_and_evaluate.py:175–224

Runs evaluation on the model using vLLM. Args: dataset: The dataset to evaluate on. vllm_rollout: The vLLM rollout object for generating responses. debug: If True, prints debug information for each sample. Returns: A dictionary containing evaluation scores: 'correct', 'partiall

(dataset, vllm_rollout, debug=True)

def evaluate_model(dataset, vllm_rollout, debug=True):
  """Runs evaluation on the model using vLLM.

  Args:
    dataset: The dataset to evaluate on.
    vllm_rollout: The vLLM rollout object for generating responses.
    debug: If True, prints debug information for each sample.

  Returns:
    A dictionary containing evaluation scores: 'correct', 'partially_correct',
    and 'correct_format' percentages.
  """
  rollout_config = base_rollout.RolloutConfig(
      max_tokens_to_generate=MAX_TOKENS_TO_GENERATE,
      max_prompt_length=MAX_PROMPT_LENGTH,
      temperature=EVALUATION_CONFIG["temperature"],
      top_p=EVALUATION_CONFIG["top_p"],
      top_k=EVALUATION_CONFIG["top_k"],
      data_type="bfloat16",
  )

  total, total_correct, total_partially_correct, total_correct_format = 0, 0, 0, 0
  for batch in tqdm(dataset):
    batch_response = vllm_rollout.generate(batch["prompt"], rollout_config)
    for i, question in enumerate(batch["question"]):
      if debug:
        print("========================================")
        print(f"Question: {question}")
        print("----------------------------------------")
        print(f"Model Generated Response: {batch_response.text[i]}")
        print("----------------------------------------")
        print(f"Target Response: {batch["target_answer"][i]}")
        print("========================================")

      is_correct, is_partially_correct, has_correct_format = score_response(
          target=batch["target_answer"][i], prediction=batch_response.text[i], debug=debug
      )
      if is_correct:
        total_correct += 1
      if is_partially_correct:
        total_partially_correct += 1
      if has_correct_format:
        total_correct_format += 1
      total += 1

  return {
      "correct": (total_correct / total) * 100,
      "partially_correct": (total_partially_correct / total) * 100,
      "correct_format": (total_correct_format / total) * 100,
  }


def safe_string_to_float(text):