shreepow/emr-medical-code-extractor

EMR Medical Code Extractor

Fine-tuned small language model for extracting ICD-10-CM, SNOMED-CT, and LOINC codes from unstructured EMR (electronic medical record) free text.

• Base model: Qwen/Qwen2.5-1.5B-Instruct
  (3.8 B params, instruction-tuned, Apache-2.0)
• Method: Supervised Fine-Tuning (SFT) with LoRA adapters
  (rank 32, alpha 64, targeting all linear layers)
• Primary dataset: rntc/mimic-icd-visit
  Real MIMIC-IV discharge summaries → ICD-10-CM diagnosis codes
• Secondary datasets:
  • JMasr/balidea-snomed-dataset – Spanish clinical text with SNOMED-CT codes
  • Tippawan/SNOMED-CT-NER-V.2 – SNOMED-CT mention-level NER
  • awacke1/LOINC-Clinical-Terminology – LOINC code dictionary
  • rntc/synthetic-loinc – Synthetic LOINC textbooks for data augmentation

---

🎯 What problem does this solve?

Clinical notes (discharge summaries, progress notes, radiology reports) contain rich information, but the structured billing / interoperability codes (ICD-10, SNOMED-CT, LOINC) are usually entered manually by trained coders. This model automates that extraction by reading the free-text narrative and generating the correct code list.

Code system	Typical use in EMR	How the model learns it
ICD-10-CM	Diagnosis billing, DRG grouping	Document-level multi-label classification from discharge summaries
SNOMED-CT	Clinical concept normalization, EHR interoperability	Entity linking from clinical mentions to concept IDs
LOINC	Lab test / observation identification	Terminology lookup from lab-mention text to LOINC codes

---

🏗️ Architecture decisions (literature-backed)

1. Generative SFT instead of multi-label classification heads
   Landmark paper MedCodER (Baksi et al., 2024) showed that decomposing coding into extraction → retrieval → reranking with an LLM outperforms vanilla classification heads on large label spaces. We adopt the same generative paradigm but train it end-to-end with SFT, which is simpler to deploy and does not require a separate retrieval index at inference time.

2. Small instruction-tuned base (1.5 B → 3.8 B effective)
   PLM-ICD (Huang et al., 2022) demonstrated that domain-specific pre-training (PubMed RoBERTa) is crucial. Rather than starting from scratch, we leverage Qwen2.5-1.5B-Instruct, which is already instruction-tuned and generalises well to medical prompts after light LoRA adaptation.

3. LoRA instead of full fine-tuning
   OpenMed-NER (Panahi, 2025) showed that lightweight domain-adaptive pre-training

   • LoRA achieves SOTA on 12 biomedical NER benchmarks while keeping training cost low. We target all linear projection layers (q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj, down_proj) with rank 32.

4. Combined multi-task training
   No single public dataset annotates all three code systems on the same documents. We therefore concatenate:

   • 60 % ICD-10-CM document-level examples (MIMIC-IV discharge summaries)
   • 20 % SNOMED-CT synthetic mention → concept pairs (from UMLS/SNOMED dictionaries)
   • 20 % LOINC synthetic mention → code pairs (from LOINC terminology tables)

   The model learns to distinguish the tasks from the prompt phrasing.

---

📁 Repository layout

├── train.py               # Single-task ICD-10-CM trainer (fastest to run)
├── train_combined.py      # Multi-code-system trainer (ICD + SNOMED + LOINC)
├── prepare_snomed_l.py    # Build synthetic SNOMED/LOINC SFT rows from terminology datasets
├── inference.py           # Run inference on a single clinical note
├── requirements.txt       # Python dependencies
└── README.md              # This file

---

🚀 Quick start – Inference (no training required)

pip install -r requirements.txt

python inference.py \
    --model_id shreepow/emr-medical-code-extractor \
    --text "Patient admitted with acute chest pain. Troponin elevated. \
            ECG shows ST elevation in leads V1-V4. Diagnosis: acute anterior \
            STEMI. PCI performed on LAD. Discharged on aspirin, clopidogrel, \
            atorvastatin and metoprolol."

Expected output:

Generated raw text:
I21.0, I21.1, Z95.5

Extracted ICD-10-CM codes:
['I21.0', 'I21.1', 'Z95.5']

To use the base model + adapter separately:

python inference.py \
    --model_id Qwen/Qwen2.5-1.5B-Instruct \
    --adapter_path shreepow/emr-medical-code-extractor \
    --text "..."

---

🏋️ Training – single task (ICD-10-CM only)

The fastest way to get a working model. Uses the MIMIC-ICD dataset with automatic fallback to a synthetic ICD-10-CM dataset if MIMIC is unavailable.

# GPU with ≥ 16 GB VRAM (T4 / L4 / A10G sufficient)
python train.py \
    --model_id Qwen/Qwen2.5-1.5B-Instruct \
    --output_dir shreepow/emr-medical-code-extractor \
    --max_train_samples 5000 \
    --max_val_samples 500 \
    --num_train_epochs 3 \
    --per_device_train_batch_size 1 \
    --gradient_accumulation_steps 8 \
    --lora_r 32 \
    --lora_alpha 64 \
    --learning_rate 2e-4 \
    --max_length 2048 \
    --eval_steps 100 \
    --save_steps 100

Key hyperparameters (based on PLM-ICD & OpenMed-NER papers):

Param	Value	Rationale
learning_rate	2e-4	Higher than full-FT because only LoRA weights update
lora_r	32	Sufficient rank for biomedical concept mapping per OpenMed-NER
lora_alpha	64	2× rank, standard heuristic
gradient_accumulation_steps	8	Effective batch size = 8 on 1 GPU
max_length	2048	Fits 1024-word discharge summaries + prompt
num_train_epochs	3	Early stopping on eval_loss prevents overfitting

The script automatically pushes the best checkpoint to the Hugging Face Hub.

---

🏋️ Training – combined multi-code system

To train a single model that handles ICD-10, SNOMED-CT and LOINC:

python train_combined.py \
    --model_id Qwen/Qwen2.5-1.5B-Instruct \
    --output_dir shreepow/emr-medical-code-extractor \
    --max_train_samples 5000 \
    --max_val_samples 500 \
    --num_train_epochs 3 \
    --lora_r 32 \
    --lora_alpha 64 \
    --learning_rate 2e-4

Data split inside the script:

• 60 % ICD-10-CM (MIMIC-ICD or synthetic fallback)
• 20 % SNOMED-CT (synthetic mention → concept from HF terminology datasets)
• 20 % LOINC (synthetic mention → code from HF LOINC datasets)

If you want to pre-build the SNOMED/LOINC synthetic rows separately:

python prepare_snomed_l.py --output_dir data/snomed_loinc_sft

Then merge them with the ICD rows in your own data pipeline before calling SFTTrainer.

---

📊 Evaluation

Both training scripts run a lightweight exact-match evaluation on the validation set after training:

• ICD-10-CM: token-level exact match (\b[A-Z]\d{2,3}\.?\d{0,3}\b)
• SNOMED-CT: digit sequence match (\b\d{6,18}\b)
• LOINC: LOINC pattern match (\b\d{3,5}-\d\b)

Metrics reported: Precision, Recall, F1 (micro-averaged per code system).

Results are saved locally as val_metrics.json and uploaded to the model repo.

---

🔬 Reproducing published baselines

If you have access to MIMIC-III/IV via PhysioNet, the gold-standard baseline is PLM-ICD (Huang et al., 2022):

• Model: RoBERTa-base-PM-M3-Voc (PubMed pre-trained)
• Method: Segment pooling (128-token chunks) + Label-Aware Attention (LAAT)
• Hyperparams: 20 epochs, lr 5e-5, linear warmup 2 k steps, batch size 8
• Expected MIMIC-III full: Micro-F1 59.8 %, Macro-F1 10.4 %

Our generative SFT approach trades a few points of raw F1 for dramatically simpler deployment (no label-vocabulary index, no attention mechanism, single forward pass).

For entity linking / normalization (SNOMED-CT, LOINC), the SOTA starting point is SapBERT (Liu et al., 2020):

• Model: cambridgeltl/SapBERT-from-PubMedBERT-fulltext
• Method: Bi-encoder with metric learning on UMLS synonymy
• Best for: ranking candidate concepts given a detected mention

You can use SapBERT as a second-pass reranker on top of our generative extractor for higher precision on rare codes.

---

🧪 Datasets used

Dataset	Task	Size	Access
rntc/mimic-icd-visit	ICD-10-CM document coding	~47 k discharge summaries	Public (MIMIC-IV credentialed subset)
FiscaAI/synth-ehr-icd10cm-prompt	Synthetic ICD-10-CM coding	~100 k samples	Public
JMasr/balidea-snomed-dataset	Spanish SNOMED-CT coding	~8 k clinical notes	Public
Tippawan/SNOMED-CT-NER-V.2	SNOMED-CT mention NER	~1 k sentences	Public
awacke1/LOINC-Clinical-Terminology	LOINC dictionary	~100 k entries	Public
rntc/synthetic-loinc	Synthetic LOINC textbooks	~50 k paragraphs	Public

---

🛠️ Extending the model

Add your own EMR data

Convert your internal notes into the same conversational JSON format:

{
  "messages": [
    {"role": "system", "content": "You are an expert medical coding assistant."},
    {"role": "user", "content": "Discharge Summary:\n...\nICD-10-CM codes (comma-separated):"},
    {"role": "assistant", "content": "I21.0, I10, E11.9"}
  ],
  "task": "icd10"
}

Then concatenate with the existing train set and resume training with --resume_from_checkpoint.

Improve SNOMED / LOINC accuracy

The current synthetic SNOMED/LOINC rows are terminology lookups (mention → code). For real-world EMR text, add mention-detection NER data (e.g. n2c2 tracks, BC5CDR, NCBI Disease) and train a two-stage pipeline:

1. Stage 1 (NER): microsoft/BiomedNLP-PubMedBERT-base-uncased fine-tuned for BIO tagging of conditions, procedures, drugs, lab tests.
2. Stage 2 (Linking): Use this generative model (or SapBERT) to map each detected span to the correct SNOMED-CT / LOINC concept.

This two-stage design is the best-supported approach in the clinical-NLP literature (Clinical NER Benchmark, 2024; OpenMed-NER, 2025).

Scale up

If you have A100 / H100 GPUs, increase the base model to:

• microsoft/Phi-3-mini-4k-instruct (3.8 B, strong clinical reasoning)
• yikuan8/Clinical-Longformer (149 M, 4 096 tokens, true long-document encoding)
• emilyalsentzer/Bio_ClinicalBERT (110 M, BERT-based, excellent for NER as stage 1)

For 7 B+ models, switch to DeepSpeed ZeRO-3 or FSDP and increase lora_r to 64-128.

---

📜 Citation

If you use this model or code, please cite the key papers that informed the design:

@article{huang2022plmicd,
  title={PLM-ICD: Automatic ICD Coding with Pretrained Language Models},
  author={Huang, Kexin and Altosaar, Jaan and Ranganath, Rajesh},
  journal={arXiv preprint arXiv:2207.05289},
  year={2022}
}

@article{baksi2024medcoder,
  title={MedCodER: A Generative AI Assistant for Medical Coding},
  author={Baksi, Ankush and others},
  journal={arXiv preprint arXiv:2409.15368},
  year={2024}
}

@article{liu2020sapbert,
  title={Self-Alignment Pretraining for Biomedical Entity Representations},
  author={Liu, Fangyu and others},
  journal={arXiv preprint arXiv:2010.11784},
  year={2020}
}

@article{panahi2025openmedner,
  title={OpenMed NER: Open-Source, Domain-Adapted State-of-the-Art Transformers for Biomedical NER Across 12 Public Datasets},
  author={Panahi, A.},
  journal={arXiv preprint arXiv:2508.01630},
  year={2025}
}

---

📄 License

• Model weights: Apache-2.0 (inherited from Qwen2.5-1.5B-Instruct)
• Training code: MIT
• Data usage: Follow the licenses of each upstream dataset (MIMIC-IV requires credentialed access via PhysioNet; synthetic datasets are CC-BY or MIT).

---

🤝 Acknowledgements

Built with 🤗 Transformers, 🤗 TRL, 🤗 PEFT, and the generous open-source clinical-NLP community. MIMIC data is provided by the MIT Lab for Computational Physiology.