benchmark/create_benchmark.py

#!/usr/bin/env python3
"""
Medical X-ray Question Generation Benchmark aka ChestAgentBench

This script generates clinical questions from X-ray case data of Eurorad dataset using GPT-4o.
It structures questions across different analytical categories and saves them as JSON.
"""

import os
import re
import json
from typing import *
from pprint import pprint

import openai
import numpy as np
from scipy import stats
import plotly.graph_objects as go
from tqdm import tqdm

from benchmark.utils import load_eurorad_dataset
from benchmark.llm import get_llm_response

# Constants
DATA_DIR = "set your data directory here, e.g. /home/MedRAX/data"
DATASET_PATH = os.path.join(DATA_DIR, "eurorad_metadata.json")

SYSTEM_PROMPT = """
You are an expert medical benchmark creation assistant.
Your goal is to generate questions that evaluate a multimodal medical AI agent's ability to interpret and reason about chest X-rays.
""".strip()

CATEGORIES_META = {
    "detection": "Identify and locate specific findings in the chest X-ray.",
    "classification": "Determine whether specific findings are present or absent in the chest X-ray.",
    "enumeration": "Count the number of target findings in the chest X-ray.",
    "localization": "Locate a given finding in the chest X-ray.",
    "comparison": "Compare the size or position of a specific finding in the chest X-ray.",
    "relationship": "Determine the relationship between two or more findings in the chest X-ray.",
    "diagnosis": "Make a diagnosis or determine a treatment plan by interpreting the chest X-ray.",
    "characterization": "Describe specific attributes (shape, density, margins, etc.) of findings.",
    "reasoning": "Explain the medical rationale and thought process behind findings and conclusions.",
}
CATEGORIES = list(CATEGORIES_META.keys())

CATEGORY_COMBINATIONS = [
    ["detection", "localization", "characterization", "reasoning"],  # Detailed Finding Analysis
    ["detection", "classification", "relationship", "reasoning"],  # Pattern Recognition & Relations
    ["localization", "comparison", "relationship", "reasoning"],  # Spatial Understanding
    ["classification", "comparison", "diagnosis", "reasoning"],  # Clinical Decision Making
    ["classification", "characterization", "diagnosis", "reasoning"],  # Diagnostic Characterization
]

DEFAULT_SECTIONS = [
    "history",
    "image_finding",
    "discussion",
    "differential_diagnosis",
    "diagnosis",
    "figures",
]


class Question:
    """A class to generate clinical questions from case data.

    This class handles creating structured clinical questions by combining case data with
    specified categories and difficulty levels.

    Attributes:
        type (str): The type of question (e.g. multiple choice)
        difficulty (str): Difficulty level of the question
        case_data (Dict[str, Any]): Dictionary containing the clinical case data
        case_content (str): Formatted case data from selected sections
        case_id (str): Unique identifier for the case
        categories (List[str]): List of analytical categories this question tests
        sections (List[str]): Case sections to include in question
        raw_content (Optional[str]): Raw LLM response to the question prompt
        content (Optional[Dict[str, str]]): Extracted content from the raw LLM response
    """

    def __init__(
        self,
        type: str,
        difficulty: str,
        case_data: Dict[str, Any],
        categories: List[str],
        sections: List[str] = [
            "history",
            "image_finding",
            "discussion",
            "differential_diagnosis",
            "diagnosis",
            "figures",
        ],
        system_prompt: str = "You are an expert medical benchmark creation assistant.",
    ) -> None:
        self.type = type
        self.difficulty = difficulty
        self.case_data = case_data
        self.case_id = case_data["case_id"]
        self.categories = categories
        self.sections = sections
        self.system_prompt = system_prompt
        self.case_content = self.select_case_sections()
        self.raw_content: Optional[str] = None
        self.content: Optional[Dict[str, str]] = None

    def create_question_prompt(self) -> str:
        """Creates a formatted prompt for generating a clinical question.

        Returns:
            str: A structured prompt containing the question parameters and clinical data
        """
        category_descriptions = "\n".join(
            f"{category}: {desc}"
            for category, desc in CATEGORIES_META.items()
            if category in self.categories
        )

        return f"""
        You must follow these guidelines:
        1. Questions must be answerable using only context and chest X-rays.
        - Questions must explicitly mention the referenced figures
        - Questions can only reference the chest X-ray figures

        2. Questions must have unambiguous, verifiable answers, and should:
        - Challenge the agent's analytical capabilities
        - Require multi-step reasoning
        - Test ability to make precise observations
        - Evaluate capability to derive insights and findings from the chest X-ray

        3. The agent has access to tools like classification, report generation, segmentation, grounding, visual question answering, etc. Your question should be complex to require the use of such tools.


        Create a {self.difficulty} {self.type} clinical question that integrates the following:

        {category_descriptions}

        based on the following clinical case:

        {self.case_content}

        Do not use any infomration derived from the CT and MRI images. Do not provide any information and findings about the chest X-rays.
        Your question should require the agent to derive insights and findings from the chest X-ray by itself.
        Your answer should be verifiable directly in the context of the case.
        You can only use the image findings that come from the chest X-ray figures.

        Your response must follow this exact format:
        THOUGHTS: [Think about different reasoning steps and tools the agent should use to answer the question]
        QUESTION: [complete question with relevant context. Incorrect choices should be very close to the correct answer.]
        FIGURES: [list of required figures, e.g. ["Figure 1", "Figure 2a"]]
        EXPLANATION: [short explanation of why your answer is verifiable in the case]
        ANSWER: [correct answer e.g. "A"]
        """.strip().replace(
            "        ", ""
        )  # remove tabs

    def select_case_sections(self) -> str:
        """Extract and format selected sections from case data into paragraphs.

        Returns:
            str: Formatted string with case sections and content
        """
        section_mapping = {
            "history": ("history", "No history provided."),
            "image_finding": ("image_finding", "No findings provided."),
            "discussion": ("discussion", "No discussion provided."),
            "differential_diagnosis": (
                "differential_diagnosis",
                "No differential diagnosis provided.",
            ),
            "diagnosis": ("diagnosis", "No diagnosis provided."),
            "figures": ("figures", "No figures provided."),
        }

        formatted = []
        for section in self.sections:
            if section in section_mapping:
                key, default = section_mapping[section]
                content = self.case_data.get(key, default)

                if key == "figures":
                    figures_text = []
                    for figure in content:
                        for subfig in figure["subfigures"]:
                            figures_text.append(f"{subfig['number']}: {subfig['caption']}")
                    content = "\n".join(figures_text)

                formatted.append(f"{section}:\n{content}")

        return "\n\n".join(formatted)

    def create_question(
        self,
        client: openai.OpenAI,
        temperature: float = 0.7,
        top_p: float = 0.95,
        max_tokens: int = 500,
        model: str = "gpt-4o",
    ) -> str:
        """Create a clinical question using LLM.

        Args:
            client (openai.OpenAI): OpenAI client instance
            temperature (float): Controls randomness in responses. Defaults to 0.7.
            top_p (float): Controls diversity via nucleus sampling. Defaults to 0.95.
            max_tokens (int): Max tokens in model response. Defaults to 500.
            model (str): OpenAI model to use. Defaults to "gpt-4o".

        Returns:
            str: LLM response containing formatted question components
        """
        self.raw_content = get_llm_response(
            client=client,
            prompt=self.create_question_prompt(),
            system_prompt=self.system_prompt,
            temperature=temperature,
            top_p=top_p,
            max_tokens=max_tokens,
            model=model,
        )
        self.content = self.extract_content()

        return self.raw_content

    def extract_content(self) -> Dict[str, str]:
        """Extract sections from raw LLM response using regex patterns.

        Returns:
            Dict[str, str]: Extracted sections including thoughts, question, figures, explanation, and answer
        """
        keywords = ["THOUGHTS", "QUESTION", "FIGURES", "EXPLANATION", "ANSWER"]

        content = {}
        for kw in keywords:
            pattern = rf"{kw}:\s*(.*?)(?=\n[A-Z]+:|$)"
            match = re.search(pattern, self.raw_content, re.DOTALL)
            content[kw.lower()] = match.group(1).strip() if match else None

        return content

    def save(self, output_path: str) -> Dict[str, Any]:
        """Save question content and metadata as a JSON file.

        Args:
            output_path (str): Directory path where the JSON file will be saved

        Returns:
            Dict[str, Any]: Question data including content (thoughts, question, figures, options,
                explanation, answer) and metadata (type, difficulty, categories, etc.)
        """
        question_metadata = self.content.copy()

        # Add metadata
        question_metadata["metadata"] = {
            "case_id": self.case_id,
            "type": self.type,
            "difficulty": self.difficulty,
            "categories": self.categories,
            "sections": self.sections,
        }

        # Create a directory for the case
        case_dir = os.path.join(output_path, str(self.case_id))
        os.makedirs(case_dir, exist_ok=True)

        # Save the question metadata to a JSON file
        output_file = os.path.join(case_dir, f"{self.case_id}_{self.__hash__()}.json")
        with open(output_file, "w") as f:
            json.dump(question_metadata, f, indent=2)

        return question_metadata


def generate_questions(
    dataset: Dict[str, Any],
    client: openai.OpenAI,
    output_dir: str,
    skip_first: int = 100,
    temperature: float = 0.7,
    top_p: float = 0.95,
    max_tokens: int = 1200,
    model: str = "gpt-4o",
) -> None:
    """Generate questions for each case and category combination.

    Args:
        dataset: Dictionary of case data
        client: OpenAI client instance
        output_dir: Directory to save generated questions
        skip_first: Number of initial cases to skip
        temperature: LLM temperature parameter
        top_p: LLM top_p parameter
        max_tokens: Maximum tokens for LLM response
        model: LLM model name
    """
    target_cases = sorted(list(dataset.keys()), key=int)[-len(dataset) : -skip_first]

    for case_id in tqdm(target_cases, desc="Processing cases"):
        case_data = dataset[case_id]

        for category in tqdm(CATEGORY_COMBINATIONS, desc=f"Categories for case {case_id}"):
            question = Question(
                type="multiple choice (A/B/C/D/E/F)",
                difficulty="complex",
                case_data=case_data,
                categories=category,
                sections=DEFAULT_SECTIONS,
                system_prompt=SYSTEM_PROMPT,
            )

            response = question.create_question(
                client=client,
                temperature=temperature,
                top_p=top_p,
                max_tokens=max_tokens,
                model=model,
            )
            question.save(output_dir)


def main():
    """Main execution function."""
    client = openai.OpenAI()

    # Load and verify dataset
    dataset = load_eurorad_dataset(
        DATASET_PATH,
        section="Chest Imaging",
        as_dict=True,
        filter_by_caption=[
            "xray",
            "x-ray",
            "x ray",
            "ray",
            "xr",
            "radiograph",
        ],
    )
    print(f"\n---\nFound {len(dataset)} cases with X-ray mentions\n---\n")

    # Optional: Print sample case for verification
    case_data = dataset["16798"]
    pprint(case_data, sort_dicts=False)

    # Generate questions
    generate_questions(dataset=dataset, client=client, output_dir="benchmark/questions")


if __name__ == "__main__":
    main()