Update README.md

README.md

@@ -20,7 +20,7 @@ The model achieves a **Rouge-2** score of 19.33 on Pubmed which is competitive t
 
 ## Usage
 
-The model can be used as follows:
+The model can be used as follows. The input is taken from the test data of the [pubmed dataset](https://huggingface.co/datasets/scientific_papers).
 
 ```python 
 LONG_ARTICLE = """"anxiety affects quality of life in those living
@@ -396,10 +396,11 @@ order to determine whether treating anxiety might
 be a potential therapy for improving fog ."""
 
 from transformers import LEDForConditionalGeneration, LEDTokenizer
+import torch
 
 tokenizer = LEDTokenizer.from_pretrained("patrickvonplaten/led-large-16384-pubmed")
 
-input_ids = tokenizer(LONG_ARTICLE).input_ids.to("cuda")
+input_ids = tokenizer(LONG_ARTICLE, return_tensors="pt").input_ids.to("cuda")
 global_attention_mask = torch.zeros_like(input_ids)
 # set global_attention_mask on first token
 global_attention_mask[:, 0] = 1