Views
No views yet
1import torch
2from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
3
4# step 1: Setup
5model_name = "StanfordAIMI/SRR-T5-SciFive"
6device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
7
8# step 2: Load Processor and Model
9model = AutoModelForSeq2SeqLM.from_pretrained(model_name).to(device)
10tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True, padding_side="right", use_fast=False)
11model.eval()
12
13# step 3: Inference (example from MIMIC-CXR dataset)
14input_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."
15inputs = tokenizer(input_text, padding="max_length", truncation=True, max_length=512, return_tensors="pt")
16inputs["attention_mask"] = inputs["input_ids"].ne(tokenizer.pad_token_id) # Add attention mask
17input_ids = inputs['input_ids'].to(device)
18attention_mask=inputs["attention_mask"].to(device)
19generated_ids = model.generate(
20 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
21 )[0]
22decoded = tokenizer.decode(generated_ids, skip_special_tokens=True)
23
24# step 4: Postprocess output
25# Remove extra <pad> tokens
26decoded = decoded.replace("<pad>", "").strip()
27
28# Split into sections based on known headers or patterns
29sections = ["History:", "Technique:", "Comparison:", "Findings:", "Impression:"]
30organs = ['Lungs and Airways:', 'Musculoskeletal and Chest Wall:','Cardiovascular:','Tubes, Catheters, and Support Devices:','Abdominal:','Pleura:','Other:','Hila and Mediastinum:']
31for section in sections:
32 decoded = decoded.replace(section, f"\n{section}")
33for organ in organs:
34 try:
35 decoded = decoded.replace(organ, f"\n{organ}")
36 except:
37 continue
38# Ensure newlines after colons and before bullet points
39decoded = decoded.replace("- ", "\n- ")
40# Ensure newlines before numbers
41for i in range(1, 8):
42 decoded = decoded.replace(f"{i}.", f"\n{i}.")
43# Remove any leading or trailing whitespace
44decoded = decoded.strip()
45print(decoded)@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}
}