SNOMED CT Entity Linking

An end-to-end pipeline for linking clinical text mentions to SNOMED CT concepts, developed in the context of the SNOMED CT Entity Linking Challenge. The pipeline detects medical entity spans in discharge summaries, links them to standardized SNOMED CT concept IDs, and optionally verifies and reranks candidates using LLM agents informed by annotation guidelines and SNOMED ontology data.

---

Background

The SNOMED CT Entity Linking Challenge provides a large corpus of de-identified clinical notes (from MIMIC-IV-Note) annotated with SNOMED CT Concept IDs. The task is to:

1. Detect mention spans in clinical text (e.g., "chest pain", "coronary artery ")
2. Link each mention to the correct SNOMED CT concept ID

The competition metric is macro-averaged character-level IoU across all SNOMED concepts.

---

Pipeline Overview

The pipeline consists of four stages:

Clinical Note
     |
     v
+--------------------------+
| Stage 1: Span Detection  |   GatorTron / GLiNER / BioClinicalBERT
|  + Dictionary Spotter    |   + Aho-Corasick dictionary matching
+--------------------------+
     |
     v
+--------------------------+
| Stage 2: Concept Linking  |   KRISSBERT / SapBERT + Qdrant vector search
+--------------------------+
     |
     v
+--------------------------+
| Stage 3a: Verification    |   LLM-based binary verification
|  (optional)               |   (Ollama / OpenAI / Anthropic / NVIDIA)
+--------------------------+
     |
     v
+--------------------------+
| Stage 3b: Reranking       |   LLM-based candidate reranking
|  (optional)               |   (Anthropic / Claude SDK / Ollama / OpenAI / NVIDIA)
+--------------------------+
     |
     v
  Predictions CSV
  (note_id, start, end, concept_id)

Stage 1 identifies medical entity spans using a fine-tuned NER model combined with an Aho-Corasick dictionary spotter. Predictions are merged and deduplicated.

Stage 2 encodes each detected mention with a biomedical language model, searches a Qdrant vector index of SNOMED CT terms, and assigns the best-matching concept ID. Candidates are enriched with FSNs, synonyms, definitions, parent concepts, and defining relationships from SNOMED.

Stage 3a (optional) uses an LLM to verify uncertain concept assignments by reasoning over the clinical context, SNOMED hierarchy, and candidate concepts.

Stage 3b (optional) uses an LLM to rerank the top-k candidates from Stage 2 using annotation guidelines, clinical context, and enriched SNOMED data embedded directly in the prompt. See Agentic Reranker README for details.

---

Obtaining Data

Three data sources are required:

1. SNOMED CT Terminology (RF2)

Register at SNOMED International and download the International Edition in RF2 format.

Place the extracted folder at:

research/data/snomedgraph/SnomedCT_InternationalRF2_PRODUCTION_20251101T120000Z/

The folder must contain Snapshot/Terminology/ with Description, Relationship, and Concept files.

2. Challenge Dataset (PhysioNet)

Access the dataset at: https://physionet.org/content/snomed-ct-entity-challenge/1.1.0/

Requires a PhysioNet account and completion of data-use certifications.

Place the files at:

research/data/mimic-iv_notes_training_set.csv
research/data/train_annotations.csv
research/data/mimic-iv_notes_test_set.csv
research/data/test_annotations.csv

3. Athena/OMOP Vocabulary

Download the SNOMED vocabulary bundle from Athena. Select at minimum the SNOMED vocabulary.

Place the files at:

research/data/athena/CONCEPT.csv
research/data/athena/CONCEPT_RELATIONSHIP.csv

---

Quick Start

Prerequisites

• Python: 3.10+ (3.12 recommended)
• GPU: NVIDIA GPU with CUDA (tested on RTX 4090, 16 GB VRAM)
• Docker: Required for Qdrant vector database
• OS: Ubuntu 24.04 LTS (tested)

1) Clone and install dependencies

git clone https://github.com/PROVIA1/snomed-ct-entity-linking-project.git
cd snomed-ct-entity-linking-project

python3 -m venv .venv
source .venv/bin/activate

pip install -r requirements.txt

Install PyTorch separately based on your CUDA version (if not already installed):

# Example for CUDA 12.x
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

2) Configure data paths

All data paths are configured in research/config.yaml. The default paths match the directory structure described in Obtaining Data. Edit if your paths differ:

data:
  snomed_rf2_dir: "research/data/snomedgraph/SnomedCT_InternationalRF2_PRODUCTION_20251101T120000Z"
  physionet:
    train_notes: "research/data/mimic-iv_notes_training_set.csv"
    train_annotations: "research/data/train_annotations.csv"
    test_notes: "research/data/mimic-iv_notes_test_set.csv"
    test_annotations: "research/data/test_annotations.csv"
  athena:
    concept_file: "research/data/athena/CONCEPT.csv"
    relationship_file: "research/data/athena/CONCEPT_RELATIONSHIP.csv"
  terms_parquet: "research/data/snomedgraph/terms.parquet"

3) Start infrastructure and build index

make start-qdrant
python3 research/scripts/run_full_pipeline.py --setup --linker-model krissbert

4) Run core services (optional but recommended for annotation workflow)

• API service docs and startup instructions: research/scripts/api/API_README.md
• Frontend UI docs and startup instructions: frontend/README.md

---

Project Structure

snomed-ct-entity-linking-project/
├── README.md
├── requirements.txt
├── research/
│   ├── config.yaml                          # Central configuration (data paths, models, infra)
│   │
│   ├── data/                                # Data directory (gitignored)
│   │   ├── snomedgraph/                     # SNOMED CT RF2 + derived files
│   │   │   ├── SnomedCT_InternationalRF2_.../
│   │   │   └── terms.parquet                # Extracted SNOMED terms (auto-generated)
│   │   ├── athena/                          # OMOP vocabulary files
│   │   ├── preprocess_data/                 # SNOMED concept mappings (auto-generated)
│   │   ├── mimic-iv_notes_training_set.csv
│   │   ├── mimic-iv_notes_test_set.csv
│   │   ├── train_annotations.csv
│   │   └── test_annotations.csv
│   │
│   ├── models/                              # Trained model checkpoints (gitignored)
│   │   ├── gliner/final_model/
│   │   ├── gatortron_lora/final_model/
│   │   └── bioclinicalbert/
│   │
│   ├── results/                             # Pipeline outputs
│   │   ├── stage1_output/
│   │   ├── stage2_output/
│   │   ├── stage3_output/
│   │   └── reranked/                        # Reranking experiment results
│   │
│   └── scripts/
│       ├── config.py                        # Configuration loader with validation
│       ├── run_full_pipeline.py              # Full pipeline orchestrator
│       │
│       ├── stage_1/                         # Stage 1: Span Detection
│       │   ├── run_stage1_pipeline.py       # Stage 1 orchestrator
│       │   ├── data.py                      # SNOMED concept mapping & dataset creation
│       │   ├── preprocess.py                # Span filtering & deduplication
│       │   ├── clean_annotations.py         # Filter inactive SNOMED concepts
│       │   ├── merge.py                     # Merge predictions from neural model and dictionary
│       │   ├── evaluate_ner_stage1_predictions.py
│       │   ├── Rule_based/
│       │   │   ├── ahocorasick_spotter.py   # Dictionary-based span detection
│       │   │   └── generate_blacklist.py    # Precision-based term filtering
│       │   ├── gatortron_lora/
│       │   │   ├── gatortron_trainer.py     # QLoRA fine-tuning
│       │   │   └── gatortron_inference.py
│       │   ├── Gliner/
│       │   │   ├── gliner_trainer.py
│       │   │   └── gliner_inference.py
│       │   └── bioclinicalbert/
│       │       ├── bioclinicalbert_trainer.py  # CRF-based training
│       │       └── bioclinicalbert_inference.py
│       │
│       ├── stage_2/                         # Stage 2: SNOMED Concept Linking
│       │   ├── run_stage2_pipeline.py       # Stage 2 orchestrator
│       │   ├── extract_snomed_terms.py      # Extract terms from RF2 to parquet
│       │   ├── krissbert_linker.py          # KRISSBERT-based linking
│       │   ├── sapbert_linker.py            # SapBERT-based linking
│       │   ├── index_qdrant_krissbert.py    # Build KRISSBERT vector index
│       │   └── index_qdrant_sapbert_meanpool.py
│       │
│       ├── concept_id_verification_agent/   # Stage 3a: LLM Verification
│       │   ├── run_verification.py          # Verification orchestrator
│       │   ├── agent.py                     # Verification agent logic
│       │   ├── tools.py                     # SNOMED lookup tools for the agent
│       │   └── __init__.py
│       │
│       ├── agentic_reranker/                # Stage 3b: LLM Reranking
│       │   ├── agent.py                     # Reranking agent + system prompt
│       │   ├── config.py                    # RerankerConfig
│       │   ├── tools.py                     # submit_reranking tool
│       │   ├── guidelines.py                # Annotation guidelines loader
│       │   ├── run_reranking.py             # Batch reranking CLI
│       │   ├── run_experiment.py            # Experiment runner with evaluation
│       │   ├── run_sdk_experiment.py        # Claude SDK + Opus + thinking
│       │   ├── run_instruction_tuning.py    # Policy extraction from prediction-vs-gold
│       │   ├── eval_diff.py                 # Before vs after evaluation
│       │   └── e2e_test.py                  # End-to-end test
│       │
│       ├── evaluation/                      # Evaluation utilities
│       │   └── report_utils.py              # IoU, span, linking, combined metrics
│       │
│       ├── eval_iou.py                      # Core IoU evaluation (competition metric)
│       ├── eval_miou.py                     # Macro IoU CLI
│       ├── eval_a2kb.py                     # A2KB (end-to-end F1) CLI
│       ├── eval_d2kb.py                     # D2KB (disambiguation precision) CLI
│       ├── eval_topk.py                     # Recall@K and MRR CLI
│       ├── evaluation_analysis.py           # Error analysis utilities
│       ├── evaluation_scenarios.py          # What-if scenario analysis
│       └── test_evaluation.py               # Unit tests for evaluation

---

Running the Pipeline

Full Pipeline

# Default: GatorTron (Stage 1) + KRISSBERT (Stage 2)
python research/scripts/run_full_pipeline.py

# Choose NER model for Stage 1
python research/scripts/run_full_pipeline.py --ner-model gliner
python research/scripts/run_full_pipeline.py --ner-model gatortron
python research/scripts/run_full_pipeline.py --ner-model bioclinicalbert

# Choose linker model for Stage 2
python research/scripts/run_full_pipeline.py --linker-model krissbert
python research/scripts/run_full_pipeline.py --linker-model sapbert

# Combine both
python research/scripts/run_full_pipeline.py --ner-model gliner --linker-model sapbert

# With Stage 3 verification
python research/scripts/run_full_pipeline.py --verify --verify-backend openai
python research/scripts/run_full_pipeline.py --verify --verify-backend ollama --verify-model qwen2.5:32b-instruct-q4_K_M

# Skip Stage 1 (reuse existing predictions)
python research/scripts/run_full_pipeline.py --skip-stage1

# Custom output path
python research/scripts/run_full_pipeline.py --output path/to/output.csv

Text Mode

Extract entities from arbitrary text and link to SNOMED (uses GLiNER for extraction):

python research/scripts/run_full_pipeline.py --text "Patient has chest pain and fever"
python research/scripts/run_full_pipeline.py --text "chest pain" --linker-model sapbert
python research/scripts/run_full_pipeline.py --text "Patient has chest pain and fever" --verify --verify-backend ollama

Stage 1 Only (Span Detection)

# Default (GatorTron)
python research/scripts/stage_1/run_stage1_pipeline.py

# Choose NER model
python research/scripts/stage_1/run_stage1_pipeline.py --ner-model gliner
python research/scripts/stage_1/run_stage1_pipeline.py --ner-model bioclinicalbert

# Text mode (GLiNER only)
python research/scripts/stage_1/run_stage1_pipeline.py --text "Patient has chest pain and fever"

Output: research/results/stage1_output/stage1_predictions.csv

Stage 2 Only (Concept Linking)

# Default (KRISSBERT)
python research/scripts/stage_2/run_stage2_pipeline.py

# With explicit input
python research/scripts/stage_2/run_stage2_pipeline.py --input research/results/stage1_output/stage1_predictions.csv

# SapBERT linker
python research/scripts/stage_2/run_stage2_pipeline.py --linker-model sapbert

# Text mode
python research/scripts/stage_2/run_stage2_pipeline.py --text "myocardial infarction" --topk 5

Output: research/results/stage2_output/linked_predictions.csv

Stage 3a (LLM Verification)

# Ollama (default, local)
python research/scripts/concept_id_verification_agent/run_verification.py --input stage2_output.csv

# OpenAI
python research/scripts/concept_id_verification_agent/run_verification.py --backend openai --input stage2_output.csv

# NVIDIA NIM (free API)
python research/scripts/concept_id_verification_agent/run_verification.py --backend nvidia --input stage2_output.csv

# Verify all predictions regardless of score
python research/scripts/concept_id_verification_agent/run_verification.py --input stage2_output.csv --verify-all

# Custom score threshold
python research/scripts/concept_id_verification_agent/run_verification.py --input stage2_output.csv --score-threshold 0.98

By default, Stage 3a only verifies predictions in the "uncertain zone" (0.96 <= score < 0.99).

Stage 3b (LLM Reranking)

Reranks top-k candidates using LLM reasoning with annotation guidelines and enriched SNOMED data. All candidate information is embedded in the prompt from the enriched Stage 2 CSV -- no external tool calls needed.

# Rerank 500 mentions with Anthropic API
python research/scripts/agentic_reranker/run_experiment.py \
    --n 500 --backend anthropic --concurrency 10

# Rerank with Opus + thinking via Claude Agent SDK (run outside Claude Code)
uv run python research/scripts/agentic_reranker/run_sdk_experiment.py \
    --n 50 --model opus --thinking-budget 10000

# Batch reranking
python research/scripts/agentic_reranker/run_reranking.py \
    --input data/2026-02-16-run-output/merged_dict_gatortron_krissbert_linked_enriched.csv \
    --backend claude_code_sdk --concurrency 5

# Batch reranking with learned policies
python research/scripts/agentic_reranker/run_reranking.py \
    --input data/2026-02-16-run-output/merged_dict_gatortron_krissbert_linked_enriched.csv \
    --backend anthropic --concurrency 10 \
    --policies research/results/reranked/tuning_policies.json

# Evaluate reranking impact
python research/scripts/agentic_reranker/eval_diff.py \
    --before original_predictions.csv \
    --after reranked_predictions.csv \
    --gold data/test-set-gold-annotations/test_annotations_cleaned.csv

See the Agentic Reranker README for full documentation.

Stage 3b Instruction Tuning (Policy Extraction)

python research/scripts/agentic_reranker/run_instruction_tuning.py \
    --predictions data/2026-02-16-run-output/merged_dict_gatortron_krissbert_linked_enriched.csv \
    --gold data/test-set-gold-annotations/test_annotations_cleaned.csv \
    --n 200 --backend anthropic --concurrency 10 \
    --output-insights research/results/reranked/tuning_insights_raw.jsonl \
    --output-policies research/results/reranked/tuning_policies.json

Setup and Indexing

# Full setup: start Qdrant, extract SNOMED terms, build vector index
python research/scripts/run_full_pipeline.py --setup

# Setup for a specific linker model
python research/scripts/run_full_pipeline.py --setup --linker-model sapbert

This will:

1. Check/start Qdrant via Docker
2. Extract SNOMED terms to terms.parquet (if missing)
3. Build the Qdrant vector index for the selected linker model (if missing)

---

Model Options

Stage 1: NER Models

Stage 3a: Verification Backends

Stage 3b: Reranking Backends

---

Evaluation

Stage 1 NER Evaluation

Evaluates span detection quality using nervaluate (strict, exact, partial, type matching):

python research/scripts/stage_1/evaluate_ner_stage1_predictions.py \
    --predictions research/results/stage1_output/stage1_predictions.csv \
    --ground_truth research/data/test_annotations_cleaned.csv \
    --notes research/data/mimic-iv_notes_test_set.csv

Macro IoU (Competition Metric)

The primary metric: macro-averaged character-level IoU across all SNOMED concepts.

python research/scripts/eval_miou.py \
    --pred research/results/stage2_output/linked_predictions.csv \
    --gold research/data/test_annotations.csv

# With score threshold filtering
python research/scripts/eval_miou.py \
    --pred research/results/stage2_output/linked_predictions.csv \
    --gold research/data/test_annotations.csv \
    --score-threshold 0.98

# Evaluate across multiple thresholds
python research/scripts/eval_miou.py \
    --pred research/results/stage2_output/linked_predictions.csv \
    --gold research/data/test_annotations.csv \
    --thresholds

Top-K Recall and MRR

Evaluates whether the correct concept appears in the top-K candidates (K=1, 5, 10, 20).

python research/scripts/eval_topk.py \
    --pred research/results/stage2_output/linked_predictions.csv \
    --gold research/data/test_annotations_cleaned.csv

Reports both end-to-end metrics (over all gold mentions) and linking-only metrics (conditioned on correct spans).

Evaluation Report

Generates a comprehensive HTML report with IoU, span, linking, and combined metrics:

python research/scripts/generate_eval_report.py \
    --pred research/results/stage2_output/linked_predictions.csv \
    --gold data/test-set-gold-annotations/test_annotations_cleaned.csv

Reranking Evaluation

Compare metrics before vs after reranking with per-mention analysis (improved/degraded/unchanged):

python research/scripts/agentic_reranker/eval_diff.py \
    --before original_predictions.csv \
    --after reranked_predictions.csv \
    --gold data/test-set-gold-annotations/test_annotations_cleaned.csv \
    --output diff_report.json