Update README.md

README.md

@@ -6,6 +6,9 @@ datasets:
 - StanfordAIMI/interpret-cxr-test-hidden
 ---
 
+<Gallery />
+
+
 # CXRMate-RRG4: Entropy-Augmented Self-Critical Sequence Training for Radiology Report Generation
 
 This is an evolution of https://huggingface.co/aehrc/cxrmate developed for the Radiology Report Generation task of BioNLP @ ACL 2024.
@@ -57,12 +60,37 @@ output_ids = model.generate(
     pixel_values=batch['images'],
     max_length=512,
     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)
 ```
 
+## Generate findings only:
+
+```python
+output_ids = model.generate(
+    pixel_values=batch['images'],
+    max_length=512,
+    num_beams=4,
+    bad_words_ids=[[tokenizer.convert_tokens_to_ids('[NF]')]],
+    eos_token_id=tokenizer.sep_token_id,
+)
+findings, _ = model.split_and_decode_sections(output_ids, tokenizer)
+```
+
+## Generate impression only:
+
+```python
+output_ids = model.generate(
+    pixel_values=batch['images'],
+    max_length=512,
+    num_beams=4,
+    bad_words_ids=[[tokenizer.convert_tokens_to_ids('[NI]')]],
+    input_ids=torch.tensor([[tokenizer.bos_token_id, tokenizer.convert_tokens_to_ids('[NF]'), tokenizer.sep_token_id]]*mbatch_size, device=device, dtype=torch.long),
+)
+_, impression = model.split_and_decode_sections(output_ids, tokenizer)
+```
+
 ## Notebook example:
 https://huggingface.co/aehrc/cxrmate-rrg24/blob/main/demo.ipynb