Views
No views yet
1import torch
2from transformers import AutoTokenizer, AutoModel
3import torch.nn.functional as F
4
5# Pick one sentence
6sentence = "The patient has a right pneumothorax."
7
8# Load pretrained model and tokenizer
9model_name = "IAMJB/RadEvalModernBERT"
10tokenizer = AutoTokenizer.from_pretrained(model_name)
11model = AutoModel.from_pretrained(model_name)
12
13# Put model in eval mode and set device
14device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
15model.to(device)
16model.eval()
17
18# Tokenize input
19inputs = tokenizer(sentence, return_tensors="pt", truncation=True, padding=True).to(device)
20
21# Get embeddings
22with torch.no_grad():
23 outputs = model(**inputs, output_hidden_states=True)
24 last_hidden_state = outputs.hidden_states[-1]
25 cls_embedding = last_hidden_state[:, 0, :] # CLS token
26 cls_embedding = F.normalize(cls_embedding, p=2, dim=1)
27
28
29print("Sentence:", sentence)
30print("Embedding shape:", cls_embedding.shape)1import argparse
2import numpy as np
3import matplotlib.pyplot as plt
4import torch
5import seaborn as sns
6from transformers import AutoTokenizer, AutoModel
7
8def get_cls_embeddings(model, tokenizer, texts, device):
9 """Get CLS token embeddings for a list of texts."""
10 embeddings = []
11
12 for text in texts:
13 # Tokenize the text
14 inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
15 inputs = {k: v.to(device) for k, v in inputs.items()}
16
17 # Get the embeddings (use CLS token)
18 with torch.no_grad():
19 outputs = model(**inputs, output_hidden_states=True)
20 # Use the last hidden state
21 last_hidden_state = outputs.hidden_states[-1]
22 # Extract CLS token (first token) embedding
23 cls_embedding = last_hidden_state[:, 0, :]
24 embeddings.append(cls_embedding.cpu().numpy()[0])
25
26 return np.array(embeddings)
27
28def compute_similarities(embeddings):
29 """Compute cosine similarity between embeddings."""
30 # Normalize embeddings
31 normalized_embeddings = embeddings / np.linalg.norm(embeddings, axis=1, keepdims=True)
32 # Compute similarity matrix
33 similarity_matrix = np.matmul(normalized_embeddings, normalized_embeddings.T)
34 return similarity_matrix
35
36def plot_heatmap(similarity_matrix, labels, output_path="cls_embedding_similarities.png"):
37 """Generate a heatmap visualization of the similarity matrix."""
38 plt.figure(figsize=(10, 8))
39
40 # Find min value to set as vmin (or use 0.6 as a reasonable value)
41 min_val = max(0.0, np.min(similarity_matrix))
42
43 # Create the heatmap with adjusted color scale
44 ax = sns.heatmap(
45 similarity_matrix,
46 annot=True,
47 fmt=".3f",
48 cmap="viridis", # Better colormap for distinguishing high values
49 vmin=min_val, # Start from minimum value or 0.6
50 vmax=1.0,
51 xticklabels=labels,
52 yticklabels=labels,
53 cbar_kws={"label": "Similarity"}
54 )
55
56 # Add title and adjust layout
57 plt.title("CLS Token Embedding Similarities")
58 plt.tight_layout()
59
60 # Rotate x-axis labels for better readability
61 plt.xticks(rotation=90)
62
63 # Save the figure
64 plt.savefig(output_path, dpi=300, bbox_inches="tight")
65 print(f"Heatmap saved to {output_path}")
66
67 # Show the plot
68 plt.show()
69
70def main():
71 # Medical terms to compare
72 medical_terms = [
73 "large right pneumothorax",
74 "right pneumothorax",
75 "pneumonia in the right lower lobe",
76 "consolidation in the right lower lobe",
77 "right 9th rib fracture",
78 "left 9th rib fracture",
79 "left 5th rib fracture",
80 "5th metatarsal fracture",
81 "no pneumothorax is present",
82 "prior consolidation has cleared",
83 "no rib fractures"
84 ]
85
86 # Set the device
87 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
88 print(f"Using device: {device}")
89
90 # Load the tokenizer
91 tokenizer = AutoTokenizer.from_pretrained(IAMJB/RadEvalModernBERT)
92
93 # Load the model
94 model = AutoModel.from_pretrained(IAMJB/RadEvalModernBERT)
95 model.to(device)
96 model.eval()
97
98 # Get CLS token embeddings for the medical terms
99 print("Generating CLS token embeddings...")
100 embeddings = get_cls_embeddings(model, tokenizer, medical_terms, device)
101
102 # Compute similarities
103 print("Computing similarity matrix...")
104 similarity_matrix = compute_similarities(embeddings)
105
106 # Plot and save the heatmap
107 print("Generating heatmap...")
108 plot_heatmap(similarity_matrix, medical_terms, "cls_embedding_similarities.png")
109
110 print("Done!")
111
112if __name__ == "__main__":
113 main()
@inproceedings{xu-etal-2025-radeval,
title = "{R}ad{E}val: A framework for radiology text evaluation",
author = "Xu, Justin and
Zhang, Xi and
Abderezaei, Javid and
Bauml, Julie and
Boodoo, Roger and
Haghighi, Fatemeh and
Ganjizadeh, Ali and
Brattain, Eric and
Van Veen, Dave and
Meng, Zaiqiao and
Eyre, David W and
Delbrouck, Jean-Benoit",
editor = {Habernal, Ivan and
Schulam, Peter and
Tiedemann, J{\"o}rg},
booktitle = "Proceedings of the 2025 Conference on Empirical Methods in Natural Language Processing: System Demonstrations",
month = nov,
year = "2025",
address = "Suzhou, China",
publisher = "Association for Computational Linguistics",
url = "https://aclanthology.org/2025.emnlp-demos.40/",
doi = "10.18653/v1/2025.emnlp-demos.40",
pages = "546--557",
ISBN = "979-8-89176-334-0",
}