Unverified Commit fb6cccb8 authored by Suraj Patil's avatar Suraj Patil Committed by GitHub
Browse files

fix qa example (#4929)

parent 38bba9cd
......@@ -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)
start_scores, end_scores = model(torch.tensor([input_ids]))
encoding = tokenizer.encode_plus(question, text, return_tensors='pt')
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])
"""
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment