Merge pull request #26 from bowang-lab/victor/benchmarking

.env.example

@@ -1,10 +0,0 @@
-OPENAI_API_KEY=
-OPENAI_BASE_URL=
-GOOGLE_API_KEY=
-GOOGLE_SEARCH_API_KEY=
-GOOGLE_SEARCH_ENGINE_ID=
-OPENROUTER_API_KEY=
-OPENROUTER_BASE_URL=
-COHERE_API_KEY=
-PINECONE_API_KEY=
-MEDGEMMA_API_URL=

benchmarking/benchmarks/rexvqa_benchmark.py

@@ -7,6 +7,8 @@ from datasets import load_dataset
 from .base import Benchmark, BenchmarkDataPoint
 from pathlib import Path
 import subprocess
+import tarfile
+import zstandard as zstd
 from huggingface_hub import hf_hub_download, list_repo_files
 
 
@@ -44,14 +46,20 @@ class ReXVQABenchmark(Benchmark):
         self.image_dataset = None
         self.image_mapping = {}  # Maps study_id to image data
         
-        super().__init__(data_dir, **kwargs)
+        # Set images_dir BEFORE parent initialization to avoid AttributeError
+        self.images_dir = f"{data_dir}/images/deid_png"
         
-        # Set images_dir after parent initialization
-        self.images_dir = f"{self.data_dir}/images/deid_png"
+        super().__init__(data_dir, **kwargs)
 
     @staticmethod
-    def download_rexgradient_images(output_dir: str = "benchmarking/data/rexvqa", repo_id: str = "rajpurkarlab/ReXGradient-160K"):
-        """Download and extract ReXGradient-160K images if not already present."""
+    def download_rexgradient_images(output_dir: str = "benchmarking/data/rexvqa", repo_id: str = "rajpurkarlab/ReXGradient-160K", test_only: bool = True):
+        """Download and extract ReXGradient-160K images if not already present.
+        
+        Args:
+            output_dir: Directory to store downloaded and extracted images
+            repo_id: HuggingFace repository ID for the dataset
+            test_only: If True, only extract images from the test split (default: True)
+        """
         output_dir = Path(output_dir)
         tar_path = output_dir / "deid_png.tar"
         images_dir = output_dir / "images"
@@ -60,6 +68,33 @@ class ReXVQABenchmark(Benchmark):
         if images_dir.exists() and any(images_dir.rglob("*.png")):
             print(f"Images already exist in {images_dir}, skipping download.")
             return
+            
+        # Load test split metadata if test_only is True
+        test_image_paths = set()
+        if test_only:
+            print("Loading test split metadata to identify test images...")
+            try:
+                # Load the test metadata to get image paths
+                test_metadata_path = output_dir / "metadata" / "test_vqa_data.json"
+                if test_metadata_path.exists():
+                    with open(test_metadata_path, 'r', encoding='utf-8') as f:
+                        test_data = json.load(f)
+                    
+                    # Extract all image paths from test data
+                    for item in test_data.values():
+                        if "ImagePath" in item and item["ImagePath"]:
+                            for rel_path in item["ImagePath"]:
+                                # Normalize path to match tar file structure
+                                norm_path = rel_path.lstrip("./")
+                                test_image_paths.add(norm_path)
+                    
+                    print(f"Found {len(test_image_paths)} test images to extract")
+                else:
+                    print("Warning: test_vqa_data.json not found, will extract all images")
+                    test_only = False
+            except Exception as e:
+                print(f"Warning: Could not load test metadata: {e}, will extract all images")
+                test_only = False
         output_dir.mkdir(parents=True, exist_ok=True)
         print(f"Output directory: {output_dir}")
         try:
@@ -96,6 +131,17 @@ class ReXVQABenchmark(Benchmark):
                                 tar_file.write(f.read())
                         else:
                             print(f"Warning: {part_file} not found, skipping...")
+                
+                # Clean up part files after successful concatenation
+                print("Cleaning up part files...")
+                for part_file in part_files:
+                    part_path = output_dir / part_file
+                    if part_path.exists():
+                        try:
+                            part_path.unlink()
+                            print(f"Deleted {part_file}")
+                        except Exception as e:
+                            print(f"Could not delete {part_file}: {e}")
             else:
                 print(f"Tar file already exists: {tar_path}")
             # Extract tar file
@@ -106,36 +152,72 @@ class ReXVQABenchmark(Benchmark):
                     print("Images already extracted.")
                 else:
                     try:
-                        subprocess.run([
-                            "tar", "-xf", str(tar_path),
-                            "-C", str(images_dir)
-                        ], check=True)
-                        print("Extraction completed!")
-                    except subprocess.CalledProcessError as e:
+                        # Stream extract with filtering for test-only images (no seeking)
+                        print("Stream extracting zstd-compressed tar file with filtering (streaming mode)...")
+
+                        # Create a decompressor
+                        dctx = zstd.ZstdDecompressor()
+
+                        # Stream extract with filtering
+                        extracted_count = 0
+                        total_png_members = 0
+
+                        with open(tar_path, 'rb') as compressed_file:
+                            with dctx.stream_reader(compressed_file) as decompressed_stream:
+                                # Use streaming tar mode to avoid seeks
+                                with tarfile.open(fileobj=decompressed_stream, mode='r|') as tar:
+                                    for member in tar:
+                                        # Only consider PNG files
+                                        if not member.isfile() or not member.name.endswith('.png'):
+                                            continue
+                                        total_png_members += 1
+
+                                        # Normalize name to match entries gathered from JSON
+                                        normalized_name = member.name.lstrip('./')
+
+                                        # Decide whether to extract this file
+                                        should_extract = True
+                                        if test_only:
+                                            should_extract = normalized_name in test_image_paths
+
+                                        if not should_extract:
+                                            # Must still advance the stream for this member
+                                            tar.members = []  # no-op in stream mode; ensure we don't hold refs
+                                            continue
+
+                                        # Ensure parent directories exist and write file by streaming
+                                        target_path = Path(images_dir) / normalized_name
+                                        target_path.parent.mkdir(parents=True, exist_ok=True)
+
+                                        extracted_file_obj = tar.extractfile(member)
+                                        if extracted_file_obj is None:
+                                            continue
+                                        with open(target_path, 'wb') as out_f:
+                                            while True:
+                                                chunk = extracted_file_obj.read(1024 * 1024)
+                                                if not chunk:
+                                                    break
+                                                out_f.write(chunk)
+
+                                        extracted_count += 1
+                                        if extracted_count % 100 == 0:
+                                            print(f"Extracted {extracted_count} test images...")
+
+                        print(f"Extraction completed! Extracted {extracted_count} matching PNGs out of {total_png_members} PNG members in the archive")
+                        
+                        # Clean up compressed tar file after successful extraction
+                        print("Cleaning up compressed tar file...")
+                        try:
+                            tar_path.unlink()
+                            print(f"Deleted {tar_path}")
+                        except Exception as e:
+                            print(f"Could not delete {tar_path}: {e}")
+                    except Exception as e:
                         print(f"Error extracting tar file: {e}")
                         return
-                    except FileNotFoundError:
-                        print("Error: 'tar' command not found. Please install tar or extract manually.")
-                        return
                 png_files = list(images_dir.rglob("*.png"))
                 print(f"Extracted {len(png_files)} PNG images to {images_dir}")
 
-                # Clean up part and tar files after successful extraction
-                print("Cleaning up part and tar files...")
-                # Remove deid_png.part* files
-                for part_file in output_dir.glob("deid_png.part*"):
-                    try:
-                        part_file.unlink()
-                        print(f"Deleted {part_file}")
-                    except Exception as e:
-                        print(f"Could not delete {part_file}: {e}")
-                # Remove deid_png.tar
-                if tar_path.exists():
-                    try:
-                        tar_path.unlink()
-                        print(f"Deleted {tar_path}")
-                    except Exception as e:
-                        print(f"Could not delete {tar_path}: {e}")
         except Exception as e:
             print(f"Error: {e}")
 
@@ -167,7 +249,7 @@ class ReXVQABenchmark(Benchmark):
         try:
             # Check for images and test_vqa_data.json, download if missing
             self.download_test_vqa_data_json(self.data_dir)
-            self.download_rexgradient_images(self.data_dir)
+            self.download_rexgradient_images(self.data_dir, test_only=True)
             
             # Construct path to the JSON file
             json_file_path = os.path.join("benchmarking", "data", "rexvqa", "metadata", "test_vqa_data.json")

benchmarking/cli.py

@@ -87,7 +87,8 @@ def run_benchmark_command(args) -> None:
         max_questions=args.max_questions,
         temperature=args.temperature,
         top_p=args.top_p,
-        max_tokens=args.max_tokens
+        max_tokens=args.max_tokens,
+        concurrency=args.concurrency
     )
     
     # Run benchmark
@@ -145,6 +146,8 @@ def main():
                            help="Maximum tokens per model response (default: 5000)")
     run_parser.add_argument("--random-seed", type=int, default=42, 
                            help="Random seed for shuffling benchmark data (enables reproducible runs, default: None)")
+    run_parser.add_argument("--concurrency", type=int, default=1,
+                            help="Number of datapoints to process in parallel (default: 1)")
     
     run_parser.set_defaults(func=run_benchmark_command)
     

benchmarking/llm_providers/base.py

@@ -85,7 +85,7 @@ class LLMProvider(ABC):
         try:
             # Simple test request
             test_request = LLMRequest(
-                text="Hello",
+                text="Hello! What model are you? Tell me your full specification.",
                 temperature=0.5,
                 max_tokens=1000
             )

benchmarking/llm_providers/medrax_provider.py

@@ -37,11 +37,11 @@ class MedRAXProvider(LLMProvider):
                 "ArcPlusClassifierTool",  # For advanced chest X-ray classification using ArcPlus
                 "ChestXRayReportGeneratorTool",  # For generating medical reports from X-rays
                 "XRayPhraseGroundingTool",  # For locating described features in X-rays
-                "MedGemmaVQATool",
-                # "XRayVQATool",  # For visual question answering on X-rays
-                # "MedicalRAGTool",  # For retrieval-augmented generation with medical knowledge
-                # "WebBrowserTool",  # For web browsing and search capabilities
-                # "DuckDuckGoSearchTool",  # For privacy-focused web search using DuckDuckGo
+                "MedGemmaVQATool", # Google MedGemma VQA tool
+                "XRayVQATool",  # For visual question answering on X-rays
+                "MedicalRAGTool",  # For retrieval-augmented generation with medical knowledge
+                "WebBrowserTool",  # For web browsing and search capabilities
+                "DuckDuckGoSearchTool",  # For privacy-focused web search using DuckDuckGo
             ]
 
             rag_config = RAGConfig(

benchmarking/runner.py

@@ -9,6 +9,7 @@ from typing import Dict, Optional, Any
 from dataclasses import dataclass
 from tqdm import tqdm
 import re
+from concurrent.futures import ThreadPoolExecutor, as_completed
 from .llm_providers import LLMProvider, LLMRequest, LLMResponse
 from .benchmarks import Benchmark, BenchmarkDataPoint
 
@@ -40,6 +41,7 @@ class BenchmarkRunConfig:
     top_p: float = 0.95
     max_tokens: int = 5000
     additional_params: Optional[Dict[str, Any]] = None
+    concurrency: int = 1
 
 
 class BenchmarkRunner:
@@ -122,40 +124,16 @@ class BenchmarkRunner:
         correct = 0
         total_duration = 0.0
         
-        # Process each data point
-        for i in tqdm(range(0, end_index), desc="Processing questions"):
+        # Determine concurrency
+        max_workers = max(1, int(getattr(self.config, "concurrency", 1) or 1))
+        
+        # Prefetch data points to avoid potential thread-safety issues inside benchmark access
+        data_points = []
+        for i in range(0, end_index):
             try:
-                data_point = benchmark.get_data_point(i)
-                
-                # Run the model on this data point
-                result = self._process_data_point(llm_provider, data_point)
-                
-                # Update counters
-                processed += 1
-                if result.is_correct:
-                    correct += 1
-                total_duration += result.duration
-                
-                # Add to results
-                self.results.append(result)
-                
-                # Save individual result immediately
-                self._save_individual_result(result)
-                
-                # Log progress
-                if processed % 10 == 0:
-                    accuracy = (correct / processed) * 100
-                    avg_duration = total_duration / processed
-                    
-                    self.logger.info(
-                        f"Progress: {processed}/{end_index} | "
-                        f"Accuracy: {accuracy:.2f}% | "
-                        f"Avg Duration: {avg_duration:.2f}s"
-                    )
-                
+                data_points.append(benchmark.get_data_point(i))
             except Exception as e:
-                self.logger.error(f"Error processing data point {i}: {e}")
-                # Add error result
+                self.logger.error(f"Error fetching data point {i}: {e}")
                 error_result = BenchmarkResult(
                     data_point_id=f"error_{i}",
                     question="",
@@ -166,10 +144,50 @@ class BenchmarkRunner:
                     error=str(e)
                 )
                 self.results.append(error_result)
-                
-                # Save individual error result immediately
                 self._save_individual_result(error_result)
-                continue
+        
+        # Process data points in parallel using a bounded thread pool
+        with tqdm(total=end_index, desc="Processing questions") as pbar:
+            with ThreadPoolExecutor(max_workers=max_workers) as executor:
+                future_to_index = {executor.submit(self._process_data_point, llm_provider, dp): idx for idx, dp in enumerate(data_points)}
+                for future in as_completed(future_to_index):
+                    idx = future_to_index[future]
+                    try:
+                        result = future.result()
+                    except Exception as e:
+                        self.logger.error(f"Error processing data point {idx}: {e}")
+                        result = BenchmarkResult(
+                            data_point_id=f"error_{idx}",
+                            question="",
+                            model_answer="",
+                            correct_answer="",
+                            is_correct=False,
+                            duration=0.0,
+                            error=str(e)
+                        )
+                    
+                    # Update counters
+                    processed += 1
+                    if result.is_correct:
+                        correct += 1
+                    total_duration += result.duration
+                    
+                    # Add to results and persist immediately
+                    self.results.append(result)
+                    self._save_individual_result(result)
+                    
+                    # Update progress bar
+                    pbar.update(1)
+                    
+                    # Periodic logging
+                    if processed % 10 == 0:
+                        accuracy = (correct / processed) * 100
+                        avg_duration = total_duration / processed if processed > 0 else 0.0
+                        self.logger.info(
+                            f"Progress: {processed}/{end_index} | "
+                            f"Accuracy: {accuracy:.2f}% | "
+                            f"Avg Duration: {avg_duration:.2f}s"
+                        )
         
         # Save final results
         summary = self._save_final_results(benchmark)
@@ -391,6 +409,7 @@ class BenchmarkRunner:
                 "model_name": self.config.model_name,
                 "benchmark_name": self.config.benchmark_name,
                 "temperature": self.config.temperature,
+                "top_p": self.config.top_p,
                 "max_tokens": self.config.max_tokens,
             },
             "benchmark_info": {

benchmarking_script.sh

@@ -0,0 +1,14 @@
+#!/bin/bash
+
+#SBATCH --job-name=chestagentbench
+#SBATCH -c 4
+#SBATCH --gres=gpu:rtx6000:1
+#SBATCH --exclude=gpu138
+#SBATCH --time=16:00:00
+#SBATCH --mem=50G
+#SBATCH --output=chestagentbench-%j.out
+#SBATCH --error=chestagentbench-%j.err
+
+source venv/bin/activate
+
+python -m benchmarking.cli run --model gpt-5 --provider medrax --system-prompt CHESTAGENTBENCH_PROMPT --benchmark chestagentbench --data-dir /scratch/ssd004/scratch/victorli/chestagentbench --output-dir temp --max-questions 2500 --concurrency 4

main.py

@@ -68,7 +68,7 @@ def initialize_agent(
     prompt = prompts[system_prompt]
 
     # Define the URL of the MedGemma FastAPI service.
-    MEDGEMMA_API_URL = os.getenv("MEDGEMMA_API_URL", "http://172.17.8.141:8002")
+    MEDGEMMA_API_URL = os.getenv("MEDGEMMA_API_URL", "http://localhost:8002")
 
     all_tools = {
         "TorchXRayVisionClassifierTool": lambda: TorchXRayVisionClassifierTool(device=device),

medgemma_script.sh

@@ -0,0 +1,14 @@
+#!/bin/bash
+
+#SBATCH --job-name=medgemma
+#SBATCH -c 4
+#SBATCH --gres=gpu:rtx6000:1
+#SBATCH --exclude=gpu138
+#SBATCH --time=16:00:00
+#SBATCH --mem=50G
+#SBATCH --output=medgemma-%j.out
+#SBATCH --error=medgemma-%j.err
+
+source medgemma/bin/activate
+
+cd medrax/tools/vqa/medgemma && python medgemma.py

medrax/docs/system_prompts.txt

@@ -17,9 +17,16 @@ Examples:
 - "Based on clinical guidelines [3], the recommended treatment approach is..."
 
 [CHESTAGENTBENCH_PROMPT]
-You are an expert medical assistant who can answer medical questions and analyze medical images with world-class accuracy.
-Use your state-of-the art reasoning and critical thinking skills to answer the questions that you are asked.
-You may use tools (if available) to complement your reasoning and you are allowed to make multiple tool calls in parallel or in sequence as needed for comprehensive answers.
-Think critically about how to best use the tools available to you and scrutinize the tool outputs.
-When encountering a multiple-choice question, your final response should end with "Final answer: \boxed{A}" from list of possible choices A, B, C, D, E, F.
-It is extremely important that you answer strictly in the format described above.
+You are a highly skilled radiology AI agent, an expert in interpreting medical images, specifically chest X-rays, CT scans, and MRIs, with world-class accuracy and precision. Your primary function is to assist in the analysis of these images and answer diagnostic questions.
+
+Your task is to provide a step-by-step, structured analysis. First, carefully examine the provided image and describe all relevant findings in a clear, concise manner. Next, use your expert medical knowledge to form a differential diagnosis based on these findings. Finally, critically evaluate the provided question and all possible choices.
+
+You have access to a suite of powerful tools to aid in your analysis. Use these tools as needed to retrieve external medical knowledge, access patient history, or perform specific image processing tasks. You should always scrutinize the output from your tools and integrate it into your reasoning. If tool outputs conflict with your initial assessment, explain the discrepancy and justify your final conclusion.
+
+Your final response for a multiple-choice question must strictly follow this format, including your step-by-step reasoning:
+1.  **Image Analysis:** [Describe image findings here]
+2.  **Differential Diagnosis:** [List possible diagnoses and their justifications]
+3.  **Critical Thinking & Tool Use:** [Show your reasoning, including how you used tools and evaluated their output]
+4.  **Final Answer:** \boxed{A}
+
+Do not provide a definitive diagnosis or treatment plan for a patient. Your purpose is to assist medical professionals with your analysis, not to replace them. You must maintain this persona and adhere to all instructions.

medrax/models/model_factory.py

@@ -123,6 +123,16 @@ class ModelFactory:
         if provider_prefix in ["openrouter"] and model_name.startswith(f"{provider_prefix}-"):
             actual_model_name = model_name[len(provider_prefix) + 1 :]
 
+        # Handle GPT-5 model
+        if model_name.startswith("gpt-5"):
+            return model_class(
+                model=actual_model_name,
+                temperature=temperature,
+                reasoning_effort="high",
+                **provider_kwargs,
+                **kwargs,
+            )
+
         # Create and return the model instance
         return model_class(
             model=actual_model_name,

medrax/tools/classification/torchxrayvision.py

@@ -12,6 +12,8 @@ from langchain_core.callbacks import (
 )
 from langchain_core.tools import BaseTool
 
+from medrax.utils.utils import preprocess_medical_image
+
 
 class TorchXRayVisionInput(BaseModel):
     """Input for TorchXRayVision chest X-ray analysis tools. Only supports JPG or PNG images."""
@@ -76,7 +78,9 @@ class TorchXRayVisionClassifierTool(BaseTool):
             ValueError: If the image cannot be properly loaded or processed.
         """
         img = skimage.io.imread(image_path)
-        img = xrv.datasets.normalize(img, 255)
+        
+        # Use robust normalization that handles both 8-bit and 16-bit images
+        img = preprocess_medical_image(img, target_range=(-1024.0, 1024.0))
 
         if len(img.shape) > 2:
             img = img[:, :, 0]

medrax/tools/segmentation/segmentation.py

@@ -20,6 +20,8 @@ from langchain_core.callbacks import (
 )
 from langchain_core.tools import BaseTool
 
+from medrax.utils.utils import preprocess_medical_image
+
 
 class ChestXRaySegmentationInput(BaseModel):
     """Input schema for the Chest X-ray Segmentation Tool."""
@@ -246,7 +248,8 @@ class ChestXRaySegmentationTool(BaseTool):
             if len(original_img.shape) > 2:
                 original_img = original_img[:, :, 0]
 
-            img = xrv.datasets.normalize(original_img, 255)
+            # Use robust normalization that handles both 8-bit and 16-bit images
+            img = preprocess_medical_image(original_img)
             img = img[None, ...]
             img = self.transform(img)
             img = torch.from_numpy(img)

medrax/utils/utils.py

@@ -1,6 +1,73 @@
 import os
 import json
-from typing import Dict, List
+import numpy as np
+from typing import Dict, List, Tuple
+
+
+def preprocess_medical_image(
+    image: np.ndarray, 
+    target_range: Tuple[float, float] = (0.0, 1.0),
+    clip_values: bool = True
+) -> np.ndarray:
+    """
+    Preprocess medical images by auto-detecting bit depth and normalizing appropriately.
+    
+    This function handles both 8-bit (0-255) and 16-bit (0-65535) images automatically,
+    normalizing them to the target range. It's designed for medical imaging tools that
+    expect consistent input ranges regardless of the original image bit depth.
+    
+    Args:
+        image (np.ndarray): Input image array (2D or 3D)
+        target_range (Tuple[float, float]): Target range for normalization (default: (0.0, 1.0))
+        clip_values (bool): Whether to clip values to target range (default: True)
+    
+    Returns:
+        np.ndarray: Normalized image in the target range
+        
+    Raises:
+        ValueError: If image is empty or has invalid values
+        ValueError: If target_range is invalid
+    """
+    if image.size == 0:
+        raise ValueError("Input image is empty")
+    
+    if len(target_range) != 2 or target_range[0] >= target_range[1]:
+        raise ValueError("target_range must be a tuple of (min, max) where min < max")
+    
+    # Convert to float for processing
+    image = image.astype(np.float32)
+    
+    # Auto-detect bit depth based on maximum value
+    max_val = np.max(image)
+    min_val = np.min(image)
+    
+    # Determine the expected maximum value based on bit depth
+    if max_val <= 255:
+        # 8-bit image
+        expected_max = 255.0
+    elif max_val <= 65535:
+        # 16-bit image
+        expected_max = 65535.0
+    else:
+        # Higher bit depth or already normalized, use actual max
+        expected_max = max_val
+    
+    # Normalize to 0-1 range first
+    if expected_max > 0:
+        image = (image - min_val) / (expected_max - min_val)
+    else:
+        # Handle edge case where image has no contrast
+        image = np.zeros_like(image)
+    
+    # Scale to target range
+    target_min, target_max = target_range
+    image = image * (target_max - target_min) + target_min
+    
+    # Clip values if requested
+    if clip_values:
+        image = np.clip(image, target_min, target_max)
+    
+    return image
 
 
 def load_prompts_from_file(file_path: str) -> Dict[str, str]: