Views
No views yet
1import torch
2from transformers import EncoderDecoderModel, AutoTokenizer
3
4# step 1: Setup
5model_name = "StanfordAIMI/SRR-BERT2BERT-RadBERT"
6device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
7
8# step 2: Load Processor and Model
9model = EncoderDecoderModel.from_pretrained(model_name).to(device)
10tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True, padding_side="right", use_fast=False)
11model.config.decoder_start_token_id = tokenizer.cls_token_id
12model.config.bos_token_id = tokenizer.cls_token_id
13model.eval()
14
15# step 3: Inference (example from MIMIC-CXR dataset)
16input_text = "CHEST RADIOGRAPH PERFORMED ON ___ COMPARISON: Prior exam from ___. CLINICAL HISTORY: Weakness, assess pneumonia. FINDINGS: Frontal and lateral views of the chest were provided. Midline sternotomy wires are again noted. The heart is poorly assessed, though remains enlarged. There are at least small bilateral pleural effusions. There may be mild interstitial edema. No pneumothorax. Bony structures are demineralized with kyphotic angulation in the lower T-spine again noted. IMPRESSION: Limited exam with small bilateral effusions, cardiomegaly, and possible mild interstitial edema."
17inputs = tokenizer(input_text, padding="max_length", truncation=True, max_length=512, return_tensors="pt")
18inputs["attention_mask"] = inputs["input_ids"].ne(tokenizer.pad_token_id) # Add attention mask
19input_ids = inputs['input_ids'].to(device)
20attention_mask=inputs["attention_mask"].to(device)
21generated_ids = model.generate(
22 input_ids, attention_mask=attention_mask, max_new_tokens=286, min_new_tokens= 120,decoder_start_token_id=model.config.decoder_start_token_id, num_beams=5, early_stopping=True, max_length=None
23 )[0]
24decoded = tokenizer.decode(generated_ids, skip_special_tokens=True)
25print(decoded)
26@article{structuring-2025,
title={Structuring Radiology Reports: Challenging LLMs with Lightweight Models},
author={Moll, Johannes and Fay, Louisa and Azhar, Asfandyar and Ostmeier, Sophie and Lueth, Tim and Gatidis, Sergios and Langlotz, Curtis and Delbrouck, Jean-Benoit},
journal={arXiv preprint arXiv:2506.00200},
url={https://arxiv.org/abs/2506.00200},
year={2025}
}