Update README.md

README.md

@@ -182,20 +182,21 @@ res = nlp(qa_input)
 # {'score': 0.984, 'start': 30, 'end': 37, 'answer': ' London'}
 
 # b) Load model & tokenizer
-model = AutoModelForQuestionAnswering.from_pretrained(model_name)
+model = AutoModelForQuestionAnswering.from_pretrained(
+  model_name, trust_remote_code=True
+)
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 
 question = f'{tokenizer.cls_token}Where do I live?'  # '<cls>Where do I live?'
 context = 'My name is Sarah and I live in London'
 encoding = tokenizer(question, context, return_tensors="pt")
-start_scores, end_scores, _, _ = model(
+output = model(
   encoding["input_ids"],
-  attention_mask=encoding["attention_mask"],
-  return_dict=False
+  attention_mask=encoding["attention_mask"]
 )
 
 all_tokens = tokenizer.convert_ids_to_tokens(encoding["input_ids"][0].tolist())
-answer_tokens = all_tokens[torch.argmax(start_scores):torch.argmax(end_scores) + 1]
+answer_tokens = all_tokens[torch.argmax(output["start_logits"]):torch.argmax(output["end_logits"]) + 1]
 answer = tokenizer.decode(tokenizer.convert_tokens_to_ids(answer_tokens))
 # 'London'
 ```