What software engineering teaches us about AI quality evaluation

import weave
from weave import EvaluationLogger
from datasets import load_dataset
from llmasajudge import LLMAsAJudge
from litellm import completion
from PIL import Image
import io
import base64
from typing import Optional

# ============================================================================
# MMMU-PRO DATASET CONSTANTS
# ============================================================================

MMMU_PRO_GRADING_PROMPT = """You are grading a multiple choice question from MMMU Pro. The correct answer is <|GROUND_TRUTH|>.

Determine if the model's response contains the correct answer.

Rules for grading:
* The model is CORRECT if they selected answer <|GROUND_TRUTH|>
* The model is INCORRECT if they selected a different answer
* Try to extract the answer from common formats like "The correct answer is (A)" or "Answer: B"
* If you cannot extract a clear answer, the model is INCORRECT

Model Response: <|RESPONSE|>

Respond with ONLY one word - either "right" or "wrong":
"""


# ============================================================================
# MMMU-PRO DATASET CLASS
# ============================================================================

class MMUProDataset:
    """MMMU Pro dataset interface"""

    def __init__(self, subset: Optional[str] = None):
        """
        Initialize MMMU Pro dataset from Hugging Face

        Args:
            subset: Subset name (default: "standard (4 options)")
        """
        print("Loading MMMU Pro dataset from Hugging Face...")

        self.subset = subset or "standard (4 options)"

        print(f"Loading subset: {self.subset}")
        self.dataset = load_dataset("MMMU/MMMU_Pro", self.subset, split='test')

        print(f"Loaded {len(self.dataset)} problems")

        # Initialize LLM judge with 'right/wrong' parser
        self.judge = LLMAsAJudge(
            models=["gpt-4o"],
            use_fully_custom_prompt=True,
            output_parser='right/wrong',
            verbose=False
        )

    def get_examples(self, num_samples: Optional[int] = None):
        """
        Returns a list of examples with questions, images (if present), and metadata

        Args:
            num_samples: Maximum number of samples to return (None = all samples)
        """
        examples = []

        dataset_size = len(self.dataset)
        samples_to_take = dataset_size if num_samples is None else min(num_samples, dataset_size)

        for i in range(samples_to_take):
            sample = self.dataset[i]

            question_text = sample.get('question', '')
            answer = sample.get('answer', '')
            options = sample.get('options', [])

            # Format question with options
            question = question_text + "\n\nChoices:\n"

            # Handle case where options might be a string instead of a list
            if isinstance(options, str):
                import ast
                try:
                    options = ast.literal_eval(options)
                except:
                    options = [options]

            for idx, option in enumerate(options):
                question += f"({chr(65 + idx)}) {option}\n"

            # Get first available image (MMMU Pro can have multiple images: image_1, image_2, etc.)
            image = None
            for img_num in range(1, 8):  # Check image_1 through image_7
                img_key = f'image_{img_num}'
                if sample.get(img_key) is not None:
                    image = sample[img_key]
                    break

            # Build grading prompt
            grading_prompt = MMMU_PRO_GRADING_PROMPT\
                .replace("<|GROUND_TRUTH|>", str(answer))\
                .replace("<|RESPONSE|>", "<RESPONSE>")

            # Build metadata
            metadata = {
                'id': sample.get('id', i),
                'topic_difficulty': sample.get('topic_difficulty', ''),
                'subject': sample.get('subject', ''),
                'img_type': sample.get('img_type', []),
                'has_image': image is not None,
                'num_options': len(options)
            }

            example = {
                'question': question,
                'answer': answer,
                'grading_prompt': grading_prompt,
                'metadata': metadata
            }

            # Add image if present
            if image is not None:
                example['image'] = image

            examples.append(example)

        return examples
    
    @weave.op 
    def get_score(self, generated_answer, grading_prompt, metadata):
        """
        Get score for a generated answer using LLMAsAJudge

        Args:
            generated_answer: The model's generated answer
            grading_prompt: The grading prompt with <RESPONSE> placeholder
            metadata: Metadata dict for reference

        Returns:
            bool: True if correct, False if incorrect
        """
        # Fill in the response
        prompt = grading_prompt.replace("<RESPONSE>", generated_answer)

        try:
            score = self.judge.judge(prompt=prompt)
            # LLMAsAJudge returns a dict: {"correct": bool, "mode": str, "votes": list}
            return score.get("correct", False)
        except Exception as e:
            print(f"Error in get_score: {e}")
            return False


# ============================================================================
# MODEL PROVIDER
# ============================================================================

def _image_to_base64(image: Image.Image) -> str:
    """Convert PIL Image to base64 string"""
    buffered = io.BytesIO()
    image.save(buffered, format="PNG")
    img_str = base64.b64encode(buffered.getvalue()).decode()
    return img_str


@weave.op
def generate_response(question: str, image: Optional[Image.Image], model_string: str) -> str:
    """
    Generate response using LiteLLM

    Args:
        question: The question text
        image: Optional PIL Image for vision tasks
        model_string: Model identifier (e.g., "gpt-4o", "anthropic/claude-3-5-sonnet-20241022")

    Returns:
        str: Model's generated answer
    """
    if image is not None:
        # Vision task with image
        image_base64 = _image_to_base64(image)

        messages = [{
            "role": "user",
            "content": [
                {"type": "text", "text": question},
                {
                    "type": "image_url",
                    "image_url": {
                        "url": f"data:image/jpeg;base64,{image_base64}"
                    }
                }
            ]
        }]
    else:
        # Text-only task
        messages = [{"role": "user", "content": question}]

    response = completion(model=model_string, messages=messages)
    answer = response["choices"][0]["message"]["content"].strip()

    return answer


# ============================================================================
# EVALUATION LOGIC
# ============================================================================

def evaluate_model(model_name: str, dataset: MMUProDataset, examples: list, project_name: str):
    """
    Evaluate a single model on MMMU-Pro dataset using original EvaluationLogger

    Args:
        model_name: Name/identifier of the model (e.g., "gpt-4o", "anthropic/claude-3-5-sonnet-20241022")
        dataset: MMUProDataset instance for scoring
        examples: List of examples to evaluate
        project_name: Weave project name

    Returns:
        dict: Evaluation results summary
    """
    print(f"\n{'='*80}")
    print(f"Evaluating Model: {model_name}")
    print(f"{'='*80}")

    # Initialize evaluation logger
    logger = EvaluationLogger(
        model=model_name,
        dataset=examples,
        name=f"mmmupro_evaluation_{model_name.replace('/', '_')}"
    )

    total_correct = 0

    # Evaluate each example
    for idx, example in enumerate(examples, 1):
        question = example["question"]
        image = example.get("image")  # May be None for text-only questions

        # Generate response
        try:
            generated_answer = generate_response(question, image, model_name)
        except Exception as e:
            print(f"✗ Example {idx}: Error generating response - {e}")
            generated_answer = "[ERROR]"

        # Score the answer
        grading_prompt = example.get("grading_prompt", "")
        metadata = example.get("metadata", {})

        is_correct = dataset.get_score(generated_answer, grading_prompt, metadata)

        if is_correct:
            total_correct += 1

        # Log prediction
        inputs_dict = {"question": question}
        if image is not None:
            inputs_dict["image"] = image

        output_dict = {"answer": generated_answer, "ground_truth": example.get("answer", "")}

        pred = logger.log_prediction(inputs=inputs_dict, output=output_dict)
        pred.log_score("correct", is_correct)
        pred.finish()

        # Print progress
        status = "✓" if is_correct else "✗"
        question_preview = question[:60].replace('\n', ' ')
        print(f"{status} Example {idx}/{len(examples)}: {question_preview}...")

    # Log summary
    accuracy = total_correct / len(examples) if len(examples) > 0 else 0.0
    print(f"\nAccuracy: {total_correct}/{len(examples)} = {accuracy:.2%}")

    logger.log_summary(
        summary={
            "total_examples": len(examples),
            "correct": total_correct,
            "accuracy": accuracy
        },
        auto_summarize=True
    )

    print(f"Evaluation complete for {model_name}!")
    print("-" * 80)

    return {
        "model": model_name,
        "total_examples": len(examples),
        "correct": total_correct,
        "accuracy": accuracy
    }


# ============================================================================
# MAIN
# ============================================================================

def main():
    """Main function to compare 2 models on MMMU-Pro dataset"""

    # Configuration
    WEAVE_PROJECT = "mmmupro-model-comparison"
    NUM_SAMPLES = 50  # Number of examples to evaluate

    # Models to compare
    MODEL_1 = "gpt-5-mini"
    MODEL_2 = "gpt-5-nano"

    print(f"MMMU-Pro Model Comparison")
    print(f"Project: {WEAVE_PROJECT}")
    print(f"Samples: {NUM_SAMPLES}")
    print(f"Model 1: {MODEL_1}")
    print(f"Model 2: {MODEL_2}")

    # Initialize Weave
    weave.init(WEAVE_PROJECT)

    # Load dataset
    dataset = MMUProDataset()
    examples = dataset.get_examples(num_samples=NUM_SAMPLES)

    print(f"\nLoaded {len(examples)} examples")

    # Evaluate both models on the same examples
    results = []

    # Evaluate Model 1
    result_1 = evaluate_model(MODEL_1, dataset, examples, WEAVE_PROJECT)
    results.append(result_1)

    # Evaluate Model 2
    result_2 = evaluate_model(MODEL_2, dataset, examples, WEAVE_PROJECT)
    results.append(result_2)

    # Print comparison summary
    print("\n" + "="*80)
    print("COMPARISON SUMMARY")
    print("="*80)
    for result in results:
        print(f"{result['model']}: {result['accuracy']:.2%} ({result['correct']}/{result['total_examples']})")
    print("="*80)
    print(f"\nView results in Weave: https://wandb.ai/your-entity/{WEAVE_PROJECT}")


if __name__ == "__main__":
    main()