Update generate_answers.py

generate_answers.py

@@ -1,15 +1,18 @@
-from transformers import BertTokenizer, BertForQuestionAnswering
+from transformers import BertForQuestionAnswering, BertTokenizer
 import torch
 
-model_path = "bert-large-uncased-whole-word-masking-finetuned-squad"
-tokenizer = BertTokenizer.from_pretrained(model_path)
-model = BertForQuestionAnswering.from_pretrained(model_path)
+model_name = "bert-large-uncased-whole-word-masking-finetuned-squad"
+tokenizer = BertTokenizer.from_pretrained(model_name)
+model = BertForQuestionAnswering.from_pretrained(model_name)
 
 def get_answer(question, context):
-    """Answers a question using BERT on given context."""
     inputs = tokenizer(question, context, return_tensors="pt", truncation=True, max_length=512)
     with torch.no_grad():
         outputs = model(**inputs)
-    start = torch.argmax(outputs.start_logits)
-    end = torch.argmax(outputs.end_logits) + 1
-    return tokenizer.convert_tokens_to_string(tokenizer.convert_ids_to_tokens(inputs["input_ids"][0][start:end]))
+
+    start_idx = torch.argmax(outputs.start_logits)
+    end_idx = torch.argmax(outputs.end_logits) + 1
+
+    answer_tokens = inputs["input_ids"][0][start_idx:end_idx]
+    answer = tokenizer.decode(answer_tokens, skip_special_tokens=True)
+    return answer