Update README.md

README.md

@@ -21,7 +21,7 @@ EAST was applied to a multimodal language model with RadGraph as the reward. Oth
  - Special tokens (`NF` and `NI`) to handle missing *findings* and *impression* sections. 
  - Non-causal attention masking for the image embeddings and a causal attention masking for the report token embeddings.
 
-## How to use:
+## Example:
 
 ```python
 import torch  
@@ -42,14 +42,23 @@ transforms = v2.Compose(
     ]
 )
 
-image = transforms(image)  # Fix.
+dataset = datasets.load_dataset('StanfordAIMI/interpret-cxr-test-public')['test']
+
+def transform_batch(batch):
+    batch['images'] = [torch.stack([transforms(j) for j in i]) for i in batch['images']]
+    batch['images'] = torch.nn.utils.rnn.pad_sequence(batch['images'], batch_first=True, padding_value=0.0)  
+    return batch
+
+dataset = dataset.with_transform(transform_batch)
+dataloader = DataLoader(dataset, batch_size=mbatch_size, shuffle=True)
+batch = next(iter(dataloader))
 
 output_ids = model.generate(
-    pixel_values=images,  # Fix.
+    pixel_values=batch['images'],
     max_length=512,
-    bad_words_ids=[[tokenizer.convert_tokens_to_ids('[NF]')], [tokenizer.convert_tokens_to_ids('[NI]')]],
     num_beams=4,
     use_cache=True,
+    bad_words_ids=[[tokenizer.convert_tokens_to_ids('[NF]')], [tokenizer.convert_tokens_to_ids('[NI]')]],
 )
 findings, impression = model.split_and_decode_sections(output_ids, tokenizer)
 ```