Update README.md

README.md

@@ -21,11 +21,10 @@ This model was used as a baseline in [Hierarchical multimodal transformers for M
 
 
 ## How to use
-
-Here is how to use this model to get the features of a given text in PyTorch:
+### Inference
+How to use this model to perform inference on a sample question and context in PyTorch:
 
 ```python
-import torch
 from transformers import LongformerTokenizerFast, LongformerForQuestionAnswering
 
 tokenizer = LongformerTokenizerFast.from_pretrained("rubentito/longformer-base-mpdocvqa")
@@ -33,17 +32,16 @@ model = LongformerForQuestionAnswering.from_pretrained("rubentito/longformer-bas
 
 text = "Huggingface has democratized NLP. Huge thanks to Huggingface for this."
 question = "What has Huggingface done?"
-encoding = tokenizer(question, text, return_tensors="pt")
-input_ids = encoding["input_ids"]
 
-# default is local attention everywhere
-# the forward method will automatically set global attention on question tokens attention_mask=encoding["attention_mask"]
+encoding = tokenizer(question, text, return_tensors="pt")
+output = model(encoding["input_ids"], attention_mask=encoding["attention_mask"])
 
-start_scores, end_scores = model(input_ids, attention_mask=attention_mask)
-all_tokens = tokenizer.convert_ids_to_tokens(input_ids[0].tolist())
+start_pos = torch.argmax(output.start_logits, dim=-1).item()
+end_pos = torch.argmax(output.end_logits, dim=-1).item()
 
-answer_tokens = all_tokens[torch.argmax(start_scores) :torch.argmax(end_scores)+1]
-answer = tokenizer.decode(tokenizer.convert_tokens_to_ids(answer_tokens))
+context_tokens = tokenizer.convert_ids_to_tokens(encoding["input_ids"][0].tolist())
+answer_tokens = context_tokens[start_pos: end_pos + 1]
+pred_answer = tokenizer.decode(tokenizer.convert_tokens_to_ids(answer_tokens))
 ```
 
 ## Model results