Dataset Loaders & Benchmarks
Relevant source files
The following files were used as context for generating this wiki page:
Dataset Loaders and Benchmarks provide a standardized interface for loading medical Question-Answering (QA) datasets into the Agentic GraphRAG system for evaluation and training. The system supports a wide range of medical benchmarks, categorized by their evaluation format, including Multiple Choice Questions (MCQ), rubric-scored conversations, and open-ended reasoning tasks.
These components reside in am_diag/loaders/dataset/ and follow a consistent architectural pattern: a base QADataset class defines the requirements for specific loaders, which then implement logic to fetch and transform raw data (primarily from HuggingFace) into structured domain models like MCQSample or RubricSample.
Sources: README.md:330-332, CONTRIBUTING.md:65-67, am_diag/loaders/dataset/init.py
Architecture and Data Flow
The dataset loading system is designed around inheritance and registration. Every loader subclasses QADataset and must implement a load() method that returns a list of QASample objects. These loaders are registered in the DATASET_REGISTRY located in am_diag/pipelines/qa.py to be accessible via CLI tools.
Component Relationship
The following diagram illustrates the relationship between the base dataset class and specialized medical benchmark implementations.
A visual representation of the class hierarchy for dataset loaders.
Sources: CONTRIBUTING.md:65-67, am_diag/loaders/dataset/init.py
Ingestion Flow
When a dataset is requested, the system streams rows from the source repository, filters invalid entries (e.g., missing questions or answers), and maps the raw fields to a unified schema.
The data flow from a CLI request to a structured list of medical samples.
Sources: tests/am_diag/loaders/dataset/test_medmcqa.py:46-55, tests/am_diag/loaders/dataset/test_pubmedqa.py:44-50
Supported Benchmark Formats
The project categorizes medical benchmarks into three primary formats, each using specific data models to represent the clinical evidence and ground truth.
Multiple Choice Questions (MCQ)
MCQ loaders process datasets where the model must select the correct answer from a set of options. Common processing steps include mapping numeric answer indices (e.g., 1, 2, 3) or text labels (e.g., yes/no/maybe) to standardized letter keys (A, B, C).
| Feature | Description | Relevant Dataset |
|---|---|---|
| Option Mapping | Converting various formats to "A", "B", "C", etc. | MedMCQADataset |
| Context Formatting | Embedding abstracts or clinical vignettes into the question stem. | PubMedQADataset |
| Metadata Extraction | Retaining medical task, body system, or difficulty level. | MedXpertQADataset |
Sources: tests/am_diag/loaders/dataset/test_medmcqa.py:126-134, tests/am_diag/loaders/dataset/test_pubmedqa.py:101-110, tests/am_diag/loaders/dataset/test_medxpertqa.py:99-110
Rubric-Scored Samples
Used for realistic health conversations where responses are evaluated against multi-turn criteria. The HealthBenchDataset is a primary example, supporting variants like "regular", "hard", and "consensus".
- Criteria Tracking: Each sample includes a list of rubrics with specific axes (accuracy, safety) and levels (example, cluster).
- Consensus Lookup: For certain variants, the loader looks up consensus criteria based on behavior categories and themes.
Sources: tests/am_diag/loaders/dataset/test_healthbench.py:53-75
Open-Ended and Reasoning
These datasets test the system's ability to generate final diagnoses from structured case reports or provide free-form public health answers.
Sources: README.md:358-366, am_diag/loaders/dataset/init.py
Implementation Details
Data Model Structure
Specific data models ensure type safety across different benchmarks.
| Model | Purpose | Key Fields |
|---|---|---|
MCQSample | Standard 4-10 option questions | question, options, answer_key, explanation |
RubricSample | Criteria-based evaluation | conversation, rubrics, criterion_ids, points_list |
OpenEndedSample | Diagnosis generation | question, answer, case_report |
Sources: am_diag/loaders/dataset/init.py:5-15, tests/am_diag/loaders/dataset/test_healthbench.py:116-120
Handling Option Shuffling
To prevent bias, loaders often include a shuffle_options flag. When enabled, the MCQSample stores the options_original to maintain a reference to the source data while providing randomized options to the model during evaluation.
Sources: tests/am_diag/loaders/dataset/test_pubmedqa_integration.py:82-84
Data Sanitization
Loaders implement logic to clean benchmark questions. For instance, the MedXpertQADataset strips explicit "Answer Choices:" markers from the question text if they are already represented in the options dictionary.
Sources: tests/am_diag/loaders/dataset/test_medxpertqa.py:121-125
Conclusion
The Dataset Loaders and Benchmarks module forms the evaluation backbone of the Agentic GraphRAG system. By abstracting the complexities of diverse HuggingFace repositories into a unified QADataset interface, the system enables consistent benchmarking across USMLE-style exams, biomedical research abstracts, and multi-turn clinical conversations. This structured approach ensures that evaluation results are comparable across different retrieval strategies and agentic reasoning loops.
Sources: README.md:330-366, CONTRIBUTING.md:65-67