fix qa example (#4929)

fb6cccb8 · Suraj Patil · GitHub · 38bba9cd · fb6cccb8
Unverified Commit fb6cccb8 authored Jun 18, 2020 by Suraj Patil Committed by GitHub Jun 17, 2020
Hide whitespace changes
Inline Side-by-side

Showing with 4 additions and 3 deletions

src/transformers/modeling_electra.py src/transformers/modeling_electra.py +4 -3

No files found.
--- a/src/transformers/modeling_electra.py
+++ b/src/transformers/modeling_electra.py
@@ -815,10 +815,11 @@ class ElectraForQuestionAnswering(ElectraPreTrainedModel):
        model = ElectraForQuestionAnswering.from_pretrained('google/electra-base-discriminator')
        question, text = "Who was Jim Henson?", "Jim Henson was a nice puppet"
-        input_ids = tokenizer.encode(question, text)
+        encoding = tokenizer.encode_plus(question, text, return_tensors='pt')
-        start_scores, end_scores = model(torch.tensor([input_ids]))
+        input_ids, token_type_ids = encoding['input_ids'], encoding['token_type_ids']
+        start_scores, end_scores = model(input_ids, token_type_ids=token_type_ids)
-        all_tokens = tokenizer.convert_ids_to_tokens(input_ids)
+        all_tokens = tokenizer.convert_ids_to_tokens(input_ids.squeeze(0))
        answer = ' '.join(all_tokens[torch.argmax(start_scores) : torch.argmax(end_scores)+1])
        """