Update README.md

README.md

@@ -21,10 +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:
+How to use this model to perform inference on a sample question and context in PyTorch:
 
 ```python
-from transformers import BigBirdForQuestionAnswering
+from transformers import BigBirdForQuestionAnswering, BigBirdTokenizerFast
 
 # by default its in `block_sparse` mode with num_random_blocks=3, block_size=64
 model = BigBirdForQuestionAnswering.from_pretrained("rubentito/bigbird-base-itc-mpdocvqa")
@@ -35,10 +35,20 @@ model = BigBirdForQuestionAnswering.from_pretrained("rubentito/bigbird-base-itc-
 # you can change `block_size` & `num_random_blocks` like this:
 model = BigBirdForQuestionAnswering.from_pretrained("rubentito/bigbird-base-itc-mpdocvqa", block_size=16, num_random_blocks=2)
 
+tokenizer = BigBirdTokenizerFast.from_pretrained("rubentito/bigbird-base-itc-mpdocvqa")
+
 question = "Replace me by any text you'd like."
 context = "Put some context for answering"
+
 encoded_input = tokenizer(question, context, return_tensors='pt')
 output = model(**encoded_input)
+
+start_pos = torch.argmax(output.start_logits, dim=-1).item()
+end_pos = torch.argmax(output.end_logits, dim=-1).item()
+
+context_tokens = tokenizer.convert_ids_to_tokens(encoded_input["input_ids"][0].tolist())
+answer_tokens = context_tokens[start_pos: end_pos]
+answer = tokenizer.decode(tokenizer.convert_tokens_to_ids(answer_tokens))
 ```
 
 ## Model results