MOSAIC: Multilingual, Taxonomy-Agnostic, and Computationally Efficient Radiological Report Classification

MOSAIC is a framework for efficient radiological report classification that is:

• 🌐 Multilingual: Works across different languages
• 🎯 Taxonomy-Agnostic: Adapts to various classification schemes
• ⚡ Computationally Efficient: Optimized for resource usage

You can read the paper here and cite it as

@inproceedings{schiavone-etal-2026-mosaic,
  title = {MOSAIC: A Multilingual, Taxonomy-Agnostic, and Computationally Efficient Approach for Radiological Report Classification in Low-Resource Settings },
  author = {Schiavone, Alice and Fraccaro, Marco and Pehrson, Lea Marie and Ingala, Silvia and Bonnevie, Rasmus and Nielsen, Michael Bachmann and Beliveau, Vincent and Ganz, Melanie and Elliott, Desmond},
  booktitle = {Proceedings of the 8th Workshop on Clinical Natural Language Processing (Clinical NLP) @ LREC 2026},
  month = {May},
  year = {2026},
  pages = {313--323},
  address = {Palma, Mallorca, Spain},
  publisher = {European Language Resources Association (ELRA)},
  editor = {Ben Abacha, Asma and Bethard, Steven and Bitterman, Danielle and Naumann, Tristan and Roberts, Kirk},
  doi = {10.63317/3qvne99cm779},
  abstract = {Radiology reports contain rich clinical information that can be used to train imaging models without relying on costly manual annotation. However, existing approaches face critical limitations: rule-based methods struggle with linguistic variability, supervised models require large annotated datasets, and recent LLM-based systems depend on closed-source or resource-intensive models that are unsuitable for clinical use. Moreover, current solutions are largely restricted to English and single-modality, single-taxonomy datasets. We introduce MOSAIC, a multilingual, taxonomy-agnostic, and computationally efficient approach for radiological report classification. Built on a compact open-access language model (MedGemma-4B), MOSAIC supports both zero-/few-shot prompting and lightweight fine-tuning, enabling deployment on consumer-grade GPUs. We evaluate MOSAIC across seven datasets in English, Spanish, French, and Danish, spanning multiple imaging modalities and label taxonomies. The model achieves a mean macro F1 score of 88 across five chest X-ray datasets, approaching or exceeding expert-level performance, while requiring only 24 GB of GPU memory. With data augmentation, as few as 80 annotated samples are sufficient to reach a weighted F1 score of 82 on Danish reports, enabling large-scale cohort classification with minimal human effort. Code and models are open-source, offering a practical alternative to large or proprietary LLMs in clinical settings.}
}

What’s in this repo

• mosaic/core/finetune.py: Unsloth + TRL SFT training with optional LoRA, early stopping, and WANDB logging.
• mosaic/core/inference.py: VLLM-based evaluation that scores predictions with weighted F1 and saves CSVs.
• mosaic/core/translate.py: Translate report datasets into multiple languages using VLLM.
• mosaic/core/perplexity.py: Perplexity utilities for MOSAIC and SIB-200 corpora.
• mosaic/core/preprocess_data.py: CLI wrappers that mirror the notebook preprocessing blocks.
• config/: Model, dataset, and experiment settings (see config/exp/README.md for experiment tags).
• scripts/: Thin shell wrappers to run training/evaluation with sensible defaults.

Setup

conda env create -f environment.yaml
conda activate mosaic
pip install -e .

Update config/paths.yaml (default base is /home/alice/work/mosaic) and config/datasets.yaml so the paths point to your local HuggingFace load_from_disk datasets.

Data layout

• Each dataset entry in config/datasets.yaml should point to a datasets-format directory containing train/, val/, and test/ splits.
• Splits are expected to include report (text), labels (stringified dict of finding → class), classes (list of class ids), findings (list of strings), and optional fs_examples for few-shot prompts.
• To recreate datasets from raw files, use the notebook-aligned CLI:

  python -m mosaic.core.preprocess_data --function mimic \
    --input-dir /path/to/raw/data \
    --output-dir data/mimic

  Available functions include mimic, padchest, casia, danskcxr, reflacx, and variants listed inside mosaic/core/preprocess_data.py.

Training

You can call the module directly or use scripts/run_training.sh (which activates the mosaic conda env and checks data paths).

python -m mosaic.core.finetune \
  --model_name medgemma-4b \              # key from config/models.yaml
  --config_tag m \                        # experiment folder in config/exp/
  --train_dataset_names "mimic" \         # space-separated keys from config/datasets.yaml
  --valid_dataset_names "mimic" \
  --output_dir outputs \
  --project_name None                     # set to a WANDB project name to enable logging

Outputs are saved under <output_dir>/models/<experiment_name>/ and checkpoints under <output_dir>/checkpoints/<experiment_name>/.

Evaluation and inference

Evaluation runs via VLLM and writes prediction/eval CSVs.

python -m mosaic.core.inference \
  --model_name medgemma-4b \              # same key used for training
  --zeroshot off \                        # one of: zeroshot | fewshot | ft-fewshot | off
  --train_dataset_names "mimic" \
  --test_dataset_names "mimic padchest_EN" \
  --models_folder outputs \               # where your trained models live
  --output_dir outputs/eval \
  --experiment_tag _m

zeroshot/fewshot run the base model, ft-fewshot adds few-shot prompts to a fine-tuned model, and off uses the fine-tuned checkpoints as-is.

Translate datasets

python -m mosaic.core.translate \
  --model_tag medgemma-4b \
  --dataset_name mimic \
  --models_folder outputs \
  --output_folder data/mimic_translated \
  --source_language eng_Latn \
  --target_languages "dan_Latn spa_Latn fra_Latn" \
  --train_only True

Languages and VLLM sampling defaults come from config/languages.yaml and config/vllm.yaml.

Perplexity utilities

python -m mosaic.core.perplexity \
  -m unsloth/gemma-3-4b-it \
  -d mosaic \
  -o outputs/ppl \
  --debug True

Use -d sib for the SIB-200 benchmark. Results are saved as CSVs.

License

This project is licensed under the MIT License - see the LICENSE file for details.