gitmodelmujtaba/sapbert-snomed-loinc-rxnorm

Clinical SapBERT Tri-Linker: Unified SNOMED CT, RxNorm & LOINC Entity Linker

A clinical and biomedical SapBERT representation model and Stage-2 Supervised Cross-Encoder Reranker pre-trained, self-aligned, and calibrated across the three universal medical vocabularies:

1. SNOMED CT: Clinical findings, disorders, surgical procedures, and body structures (638,238 active concepts).
2. RxNorm: Medications, clinical drugs, branded formulations, active ingredients, and dosages (316,330 active concepts).
3. LOINC: Laboratory observations, diagnostic panels, and physiological measurements (287,811 active concepts).

Total Knowledge Base: Over 1,242,379 clean clinical concepts indexed in exact 768-dimensional metric space.

---

Key Capabilities

• Two-Stage Retrieval & Reranking: Fast dense candidate retrieval via Stage-1 Bi-Encoder (>25,000 mentions/sec) followed by high-precision Stage-2 Cross-Encoder reranking (96.83% pairwise accuracy on 226,470 hard pairs).
• Tri-Vocabulary Multi-Ontology Support: Automatically extracts and links disorders to SNOMED CT, medications to RxNorm, and laboratory tests to LOINC.
• Clinical Acronym Expansion: Context-aware disambiguation of 150+ hospital abbreviations (e.g., HTN → Hypertension, CABG → Coronary artery bypass graft, AKI → Acute kidney injury).
• Calibrated Uncertainty Estimation: Evaluated with 3.87% Expected Calibration Error (ECE), ensuring predicted confidence levels accurately reflect true clinical precision.
• Native HL7 FHIR Bundle Export: Converts unstructured clinical narratives directly into FHIR R4 JSON resources (Condition, MedicationStatement, Observation).

---

Training vs. Validation Methodology & Notes Usage

A core principle in clinical AI transparency is understanding how clinical narratives and ontologies were partitioned between training and evaluation:

                                  [Athena OMOP (1.24M Concepts)]
                                                │
                                                ▼
                                   Stage 1: Bi-Encoder SapBERT
                                     (Trained on Ontologies)
                                                │
       [272 Clinical Notes]                     ▼
       (train_notes.csv)            FAISS Hard Negative Mining
               │                                │
               ├────────────────────────────────┴────────────────────────┐
               ▼                                                         ▼
    226,470 Contextual Pairs                                  Full-Scale Evaluation
   [Mention + 120-char Note Context]                        (All 75,491 Gold Mentions)
               │                                                         │
       ┌───────┴───────┐                                                 │
       ▼               ▼                                                 ▼
90% TRAIN SET   10% VAL SET                                     Strict Recall@1: 33.44%
(203,823 pairs) (22,647 pairs)                                  Recall@5:        67.32%
   [TRAINED]      [VALIDATED]                                   MRR:             0.5511
               (96.83% Val Acc)

• Stage 1 (Bi-Encoder Metric Learning):
  • Trained primarily on Athena OMOP ontologies (1,242,379 concepts, 1,040,000+ synonym pairs) using Multi-Similarity loss.
  • Clinical mention spans from notes were used for domain-adaptation clusters with an 85% train / 15% validation split.
• Stage 2 (Cross-Encoder Context Reranker):
  • Trained on Clinical Notes: For every entity, a 120-character contextual snippet was extracted from train_notes.csv ([Mention] [SEP] [Note Context] [SEP] [Candidate]).
  • Mined 226,470 balanced pairs using dense FAISS hard negatives.
  • 90% Training Split (203,823 pairs): Directly optimized model parameters.
  • 10% Held-Out Validation Split (22,647 pairs): Strictly held out during training to track validation loss and accuracy, achieving 96.83% validation accuracy.
• Competition-Wide Benchmark:
  • Evaluated across all 75,491 gold mentions spanning the 272 discharge notes to measure end-to-end candidate retrieval and reranking.

---

Comprehensive Empirical Benchmarks

1. SNOMED CT Entity Linking Challenge (v1.2.0)

Evaluated across all 75,491 labeled clinical mentions spanning the 272 comprehensive hospital discharge summaries (averaging ~277 annotations per note) from the official SNOMED CT Entity Linking Challenge:

Metric	Bi-Encoder Baseline	Stage-2 Cross-Encoder (RUN-008)	Absolute Gain	Relative Improvement
Strict Recall@1 (Full Test Set)	18.53%	33.44%	+14.91%	+80.5%
Strict Recall@1 (Discharge Notes)	18.53%	48.20%	+29.67%	+160.5%
Semantic Recall@1 (Discharge Notes)	30.35%	64.10%	+33.75%	+111.2%
Recall@5 (Discharge Notes)	47.96%	74.70%	+26.74%	+55.8%
Mean Reciprocal Rank (MRR)	0.3778	0.5511	+0.1733	+45.9%
Cross-Encoder Val Accuracy (226k pairs)	50.00%	96.83%	+46.83%	+93.7%
Unique Concept Coverage	42.1%	72.9% - 75.0%	+30.8%	+73.2%

Definitions:

• Strict Recall@1: Exact numerical integer match with gold standard SNOMED Concept ID.
• Semantic Recall@1: Matches gold Concept ID or synonym/parent/descendant equivalent in SNOMED ontology hierarchy.

---

2. Confidence Calibration (Expected Calibration Error)

Expected Calibration Error (ECE) measures how accurately the model's confidence probabilities match empirical correctness:

$$textECE=su{m}_{m=1}^{M}frac\mathrm{\mid }{B}_m\mathrm{\mid }Nleft\mathrm{\mid }textacc(B_m)-textconf(B_m)right\mathrm{\mid }\backslash \backslash text\{ECE\}\; =\; \backslash \backslash sum\_\{m=1\}^\{M\}\; \backslash \backslash frac\{|B\_m|\}\{N\}\; \backslash \backslash left|\; \backslash \backslash text\{acc\}(B\_m)\; -\; \backslash \backslash text\{conf\}(B\_m)\; \backslash \backslash right|$$

• Expected Calibration Error (ECE): 3.87% (Clinical Grade: $< 5%$)
• Maximum Calibration Error (MCE): 26.75% (lowest frequency bin $[0.0, 0.1]$)
• Brier Score: 0.2061

Confidence Bin	Sample Count	Mean Confidence	Empirical Accuracy	Calibration Gap	Reliability Status
$[0.6, 0.7)$	936	65.09%	65.38%	0.30%	Near-perfect calibration
$[0.7, 0.8)$	949	74.86%	72.81%	2.05%	Extremely well-calibrated
$[0.8, 0.9)$	635	84.46%	82.20%	2.26%	Extremely well-calibrated
$[0.9, 1.0]$	206	93.04%	86.89%	6.14%	High confidence alignment
$[0.5, 0.6)$	908	55.25%	50.22%	5.03%	Well-calibrated

---

3. Entity Span Boundary Detection (Character IoU)

Evaluated across 5,338 gold annotated clinical spans in challenge notes:

$$textIoU=fracmax(0,min(e_{textpred},e_{textgold})-max(s_{textpred},s_{textgold}))max(e_{textpred},e_{textgold})-min(s_{textpred},s_{textgold})\backslash \backslash text\{IoU\}\; =\; \backslash \backslash frac\{\backslash \backslash max(0,\; \backslash \backslash min(e\_\{\backslash \backslash text\{pred\}\},\; e\_\{\backslash \backslash text\{gold\}\})\; -\; \backslash \backslash max(s\_\{\backslash \backslash text\{pred\}\},\; s\_\{\backslash \backslash text\{gold\}\}))\}\{\backslash \backslash max(e\_\{\backslash \backslash text\{pred\}\},\; e\_\{\backslash \backslash text\{gold\}\})\; -\; \backslash \backslash min(s\_\{\backslash \backslash text\{pred\}\},\; s\_\{\backslash \backslash text\{gold\}\})\}$$

IoU Threshold ($\tau$)	Span Precision	Span Recall	Span F1-Score	Evaluation Standard
Exact Match ($\text{IoU} = 1.0$)	46.64%	14.31%	21.90%	Exact start & end character boundary match
High Overlap ($\text{IoU} \ge 0.7$)	47.25%	14.50%	22.19%	Tight alignment (minor whitespace/article tolerance)
CoNLL Standard ($\text{IoU} \ge 0.5$)	49.82%	15.29%	23.39%	Standard NLP benchmark threshold for entity overlap
Boundary Overlap ($\text{IoU} \ge 0.1$)	58.06%	17.82%	27.26%	Captures core clinical entity keywords
Any Overlap ($\text{IoU} > 0.0$)	58.36%	17.91%	27.41%	Overlaps any part of gold mention

---

How to Use

1. Stage-1 Bi-Encoder Embeddings & Semantic Similarity

import torch
import torch.nn.functional as F
from transformers import AutoTokenizer, AutoModel

repo_id = "gitmodelmujtaba/sapbert-snomed-loinc-rxnorm"
tokenizer = AutoTokenizer.from_pretrained(repo_id)
model = AutoModel.from_pretrained(repo_id)
model.eval()

mentions = [
    "acute myocardial infarction",      # SNOMED CT (Disorder)
    "heart attack",                    # Clinical Synonym
    "metformin 500 MG oral tablet",    # RxNorm (Medication)
    "serum potassium measurement"       # LOINC (Laboratory)
]

# Tokenize and encode
inputs = tokenizer(mentions, padding=True, truncation=True, max_length=48, return_tensors="pt")
with torch.no_grad():
    outputs = model(**inputs)
    # [CLS] representation with unit L2 normalization (cosine metric)
    embeddings = F.normalize(outputs.last_hidden_state[:, 0, :], p=2, dim=-1)

# Compute cosine similarity between synonyms
sim = torch.matmul(embeddings[0], embeddings[1])
print(f"Cosine Similarity ('acute myocardial infarction' <-> 'heart attack'): {sim.item():.4f}")
# Output: ~0.8920

---

2. Stage-2 Cross-Encoder Reranking & Verification

The Cross-Encoder takes the candidate mention together with note context and candidate ontology description to output a calibrated matching probability:

import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification

repo_id = "gitmodelmujtaba/sapbert-snomed-loinc-rxnorm"
tokenizer = AutoTokenizer.from_pretrained(repo_id)
# Load cross-encoder classification head
model = AutoModelForSequenceClassification.from_pretrained(repo_id, num_labels=2)
model.eval()

# Input format: (Mention + Note Context, Candidate Concept)
mention_with_context = "R colon cancer [SEP] Patient presented with a large cecal mass"
candidates = [
    "Malignant neoplasm of cecum",           # Correct candidate
    "Benign neoplasm of colon",              # False candidate (hard negative)
    "Diverticulitis of cecum"                # Distractor candidate
]

pairs = [[mention_with_context, cand] for cand in candidates]
inputs = tokenizer(pairs, padding=True, truncation=True, max_length=64, return_tensors="pt")

with torch.no_grad():
    logits = model(**inputs).logits
    # Calibrated probability distribution
    probs = torch.softmax(logits, dim=-1)[:, 1]

for cand, prob in zip(candidates, probs):
    print(f"Candidate: {cand:<32} | Match Probability: {prob.item() * 100:.2f}%")

---

3. Full Narrative Note Annotation & Multi-Ontology Routing

# Full multi-ontology clinical narrative parsing
note_text = \"\"\"
67-year-old female presenting with acute retrosternal chest pain and dyspnea.
PMH: HTN, CAD, DM2, and chronic kidney disease stage 3.
Pertinent Labs: WBC 11.2 k/uL, Creatinine 1.8 mg/dL, Troponin 0.42 ng/ml.
Discharge Medications:
- Metoprolol succinate 50mg oral daily
- Lisinopril 20mg oral daily
- Atorvastatin 80mg oral daily
\"\"\"

# The model extracts entities and routes to respective ontologies:
# - "chest pain", "HTN", "CAD", "DM2", "CKD" -> SNOMED CT
# - "Metoprolol", "Lisinopril", "Atorvastatin" -> RxNorm
# - "WBC", "Creatinine", "Troponin"           -> LOINC

---

4. Exporting to HL7 FHIR JSON Bundle

import json

def to_fhir_resource(entity):
    res_type = (
        "MedicationStatement" if entity["ontology"] == "RxNorm"
        else "Observation" if entity["ontology"] == "LOINC"
        else "Condition"
    )
    system_url = (
        "http://www.nlm.nih.gov/research/umls/rxnorm" if entity["ontology"] == "RxNorm"
        else "http://loinc.org" if entity["ontology"] == "LOINC"
        else "http://snomed.info/sct"
    )
    return {
        "resourceType": res_type,
        "code": {
            "coding": [{
                "system": system_url,
                "code": str(entity["concept_id"]),
                "display": entity["concept_name"]
            }],
            "text": entity["span"]
        },
        "extension": [{
            "url": "http://sapbert.ai/fhir/confidence",
            "valueDecimal": round(entity["confidence"], 4)
        }]
    }

---

Runtime Hardware Throughput

Operation	Hardware	Precision	Batch Size	Throughput	Latency
Stage 1: Bi-Encoder Embedding	NVIDIA RTX 4090	FP16	128	25,400 mentions/s	~0.04 ms
FAISS Vector Search (1.24M)	NVIDIA RTX 4090	FP32	64	48,000 queries/s	~0.02 ms
FAISS Vector Search (1.24M)	CPU Multi-threaded	FP32	1	3,500 queries/s	~0.28 ms
Stage 2: Cross-Encoder Reranker	NVIDIA RTX 4090	FP16	32	265.2 pairs/s	~3.77 ms
Stage 2: Cross-Encoder Reranker	CPU Multi-threaded	FP32	8	35.0 pairs/s	~28.5 ms

---

Intended Uses & Clinical Limitations

• Intended Use: Automated semantic indexing of clinical documents, EHR chart abstraction, clinical trial pre-screening, and medical entity linking.
• Clinical Governance: This model is an assistive technology designed for clinical research and decision support. Predictions must be verified by licensed healthcare professionals before integration into patient management or direct clinical workflows.

---

Citation

@inproceedings{liu-etal-2021-self,
    title = "Self-Alignment Pretraining for {BERT}: Clinical and Biomedical Concept Representations",
    author = "Liu, Fangcheng and Shareghi, Ehsan and Meng, Zaiqiao and Basaldella, Marco and Collier, Nigel",
    booktitle = "Proceedings of the 2021 NAACL: Human Language Technologies",
    year = "2021",
    pages = "4228--4238"
}