Yale-BIDS-Chen / medpmc-screening-pubmedbert-caption-reference

MedPMC Initial Screening Model

This model is used in the initial screening stage of the MedPMC data curation pipeline. It is a text-based classifier that takes a figure caption and its inline reference text from a PubMed Central article as input and predicts whether the corresponding figure is likely to be a clinically relevant medical image for downstream multimodal data curation.

The model is initialized from microsoft/BiomedNLP-BiomedBERT-base-uncased-abstract-fulltext .

Task

The model performs binary classification.

Input format

This repository corresponds to the caption + reference text version of the initial screening model. The input text should concatenate the figure caption and inline reference text using the following format:

"Caption": {figure_caption}
"Reference Text": {reference_text_1}
{reference_text_2}
...

Quick start

import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification

repo_id = "Yale-BIDS-Chen/medpmc-screening-pubmedbert-caption-reference"

tokenizer = AutoTokenizer.from_pretrained(repo_id)
model = AutoModelForSequenceClassification.from_pretrained(repo_id)
model.eval()

caption = "Axial CT image showing a pulmonary nodule in the right upper lobe."
references = [
    "The CT findings demonstrated a solitary pulmonary nodule.",
    "Follow-up imaging was recommended."
]

text = '"Caption": ' + caption + '\n"Reference Text": ' + "\n".join(references)

inputs = tokenizer(
    text,
    return_tensors="pt",
    truncation=True,
    padding=True,
    max_length=512,
)

with torch.no_grad():
    outputs = model(**inputs)
    probs = torch.softmax(outputs.logits, dim=-1)
    pred = torch.argmax(probs, dim=-1).item()

print("Prediction:", pred)
print("Probabilities:", probs.tolist())

Model Performance

MedPMC includes multiple initial screening variants depending on the input text and model backbone. The table below summarizes the performance of different screening models evaluated on the MedPMC validation set.