StanfordAIMI / SRR-BERT-Upper

Usage:

import json
import torch
from transformers import BertTokenizer, BertForSequenceClassification
from datasets import load_dataset
import requests

# Configuration
MODEL_PATH = "StanfordAIMI/SRR-BERT-Upper"
MAPPING_URL = "https://raw.githubusercontent.com/jbdel/StructEval/refs/heads/main/structeval/upper_mapping.json"
MAX_LENGTH = 128
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Fetch mapping from GitHub
resp = requests.get(MAPPING_URL)
resp.raise_for_status()
label_map = resp.json()
idx2label = {v: k for k, v in label_map.items()}

# Load tokenizer & model
tokenizer = BertTokenizer.from_pretrained("microsoft/BiomedVLP-CXR-BERT-general")
model = BertForSequenceClassification.from_pretrained(MODEL_PATH, num_labels=len(label_map))
model.to(DEVICE).eval()

# Grab one test sentence
dataset = load_dataset("StanfordAIMI/StructUtterances", split="test_reviewed")
sentence = dataset[35]["utterance"]

# Tokenize and infer
inputs = tokenizer(
    sentence,
    padding="max_length",
    truncation=True,
    max_length=MAX_LENGTH,
    return_tensors="pt"
).to(DEVICE)

with torch.no_grad():
    logits = model(**inputs).logits
    preds = (torch.sigmoid(logits)[0].cpu().numpy() > 0.5).astype(int)

pred_labels = [idx2label[i] for i, flag in enumerate(preds) if flag]

print(f"Sentence: {sentence}")
print("Predicted labels:", pred_labels)

Output:

Sentence: Patchy consolidation in the left retrocardiac area, suggestive of atelectasis or early airspace disease.
Predicted labels: ['Consolidation', 'Air space opacity']

Label Mapping:

{
    "Pleural Effusion": 0,
    "Upper abdominal finding": 1,
    "Widened cardiac silhouette": 2,
    "Lung Finding": 3,
    "No Finding": 4,
    "Widened aortic contour": 5,
    "Pleural Thickening": 6,
    "Vascular finding": 7,
    "Consolidation": 8,
    "Pneumothorax": 9,
    "Subdiaphragmatic gas": 10,
    "Masslike opacity": 11,
    "Chest wall finding": 12,
    "Focal air space opacity": 13,
    "Segmental collapse": 14,
    "Fracture": 15,
    "Mediastinal mass": 16,
    "Solitary masslike opacity": 17,
    "Support Devices": 18,
    "Mediastinal finding": 19,
    "Pleural finding": 20,
    "Air space opacity": 21,
    "Diffuse air space opacity": 22,
    "Multiple masslike opacities": 23,
    "Musculoskeletal finding": 24
  }

Classification Report:

                             precision    recall  f1-score   support

           Pleural Effusion       0.97      0.98      0.98     17162
    Upper abdominal finding       0.00      0.00      0.00         0
 Widened cardiac silhouette       0.97      0.96      0.96      5596
               Lung Finding       0.87      0.88      0.87      9305
                 No Finding       0.93      0.91      0.92     59962
     Widened aortic contour       0.90      0.98      0.94      1782
         Pleural Thickening       0.86      0.94      0.90      3098
           Vascular finding       0.95      0.91      0.93      2021
              Consolidation       0.95      0.98      0.97     27234
               Pneumothorax       0.87      0.92      0.90      5535
       Subdiaphragmatic gas       0.96      0.88      0.92       342
           Masslike opacity       0.00      0.00      0.00         0
         Chest wall finding       0.98      0.99      0.99      1115
    Focal air space opacity       0.87      0.62      0.72      1442
         Segmental collapse       0.91      0.86      0.88      1097
                   Fracture       0.85      0.93      0.89      3213
           Mediastinal mass       0.79      0.87      0.83       573
  Solitary masslike opacity       0.93      0.95      0.94      4056
            Support Devices       0.74      0.70      0.72      9181
        Mediastinal finding       0.88      0.94      0.91      1389
            Pleural finding       0.86      0.82      0.84       771
          Air space opacity       0.80      0.74      0.77      4816
  Diffuse air space opacity       0.96      0.98      0.97     10049
Multiple masslike opacities       0.92      0.84      0.88        55
    Musculoskeletal finding       0.83      0.82      0.83        55

                  micro avg       0.92      0.92      0.92    169849
                  macro avg       0.82      0.82      0.82    169849
               weighted avg       0.92      0.92      0.92    169849
                samples avg       0.92      0.92      0.91    169849