Adding Custom Loaders

Relevant source files

The following files were used as context for generating this wiki page:

• CONTRIBUTING.md
• am_diag/pipelines/ingestion.py
• am_diag/chunking/markitdown.py
• am_diag/loaders/corpus/statpearls.py
• am_diag/loaders/dataset/medmcqa.py
• README.md

Custom loaders are the entry points for data into the Agentic GraphRAG system. They are categorized into two types: Corpus Loaders, which stream raw medical text into the Ingestion Pipeline for graph construction, and Dataset Loaders, which provide structured questions and answers for evaluation via the Retrieval Pipeline.

Sources: CONTRIBUTING.md:76-80, README.md:120-130

Corpus Loaders

Corpus loaders are responsible for streaming medical literature from external sources (like HuggingFace) and converting them into standardized Document objects. These loaders utilize asynchronous streaming to ensure that large-scale datasets do not block the event loop.

Implementation Requirements

To add a new corpus loader, you must:

1. Subclass CorpusLoader: Located in am_diag/loaders/corpus/.
2. Define Class Variables: Set corpus_name and hf_repo.
3. Implement _row_to_document: A method that transforms a raw data row into a Document object.
4. Register the Loader: Add the new class to the CORPUS_LOADER_REGISTRY in am_diag/pipelines/ingestion.py.

Sources: CONTRIBUTING.md:77-79, am_diag/pipelines/ingestion.py:34-40

Ingestion Flow for Corpus Data

The following diagram illustrates how a corpus loader interacts with the ingestion stack to populate the Knowledge Graph (KG) and Vector Store.

The system uses asyncio.to_thread for HuggingFace iterators to prevent blocking.
Sources: CONTRIBUTING.md:78-79, am_diag/pipelines/ingestion.py:90-100

Dataset Loaders

Dataset loaders are used for benchmarking and evaluation. They load structured samples (e.g., Multiple Choice Questions) into the QA system.

Implementation Requirements

1. Subclass QADataset: Located in am_diag/loaders/dataset/.
2. Define Class Variables: Set dataset_name and hf_repo.
3. Implement load(): This method should return a list[QASample].
4. Register the Loader: Add the new class to the DATASET_REGISTRY in am_diag.pipelines.qa.

Sources: CONTRIBUTING.md:80-81, am_diag/loaders/dataset/medmcqa.py:25-30

Dataset Support Matrix

The project supports various medical benchmark formats:

Format	Example Datasets	Description
MCQ (Exact-Match)	MedQA, MedMCQA, MMLU-Med	4-10 option multiple choice questions.
Rubric-Scored	HealthBench, RAR-Med	Evaluated against specific clinical criteria.
Open-Ended	MedCaseReasoning, CareQA	Diagnosis generation from case reports.

Sources: README.md:235-275

Specialized Loaders: MarkItDown

For non-HuggingFace files (PDF, DOCX, HTML, Excel), the system provides the MarkItDownLoader. This loader wraps the markitdown library to convert local files into Document objects with full provenance metadata.

Configuration Options

The MarkItDownLoader supports advanced initialization parameters for complex document processing:

Parameter	Type	Description
llm_client	Any	OpenAI-compatible client for generating image descriptions.
docintel_endpoint	str	Azure Document Intelligence endpoint for OCR.
enable_plugins	bool	Enables third-party MarkItDown plugins.
requests_session	Any	Custom session for HTTP-based fetching.

Sources: am_diag/chunking/markitdown.py:24-42

Provenance Metadata

When MarkItDownLoader converts a file, it automatically populates the properties field of the Document with metadata:

• file_path: Absolute path to the source file.
• file_name: Name including extension.
• file_extension: The file suffix.
• file_size_bytes: Size on disk.
• loader: Set to "markitdown".

Sources: am_diag/chunking/markitdown.py:100-110

Summary of Registration

To ensure the system recognizes custom loaders, they must be registered in the appropriate registry within the pipelines module.

# am_diag/pipelines/ingestion.py

CORPUS_LOADER_REGISTRY: dict[str, type[CorpusLoader]] = {
    "clinical_guidelines": ClinicalGuidelinesCorpusLoader,
    "pubmed": PubMedCorpusLoader,
    "statpearls": StatPearlsCorpusLoader,
    # Add your custom loader here
}

Sources: am_diag/pipelines/ingestion.py:34-40

Adding custom loaders allows the system to expand its medical knowledge base and evaluation capabilities by integrating new literature and specialized medical benchmarks.