Update app.py

app.py

@@ -1,32 +1,35 @@
 import torch
 import streamlit as st
-from transformers import BertTokenizer, BertModel
+from transformers import BertTokenizer, BertForQuestionAnswering
 
-# Load the pre-trained model and tokenizer
+# Utilize BertForQuestionAnswering model for direct start/end logits
 tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
-model = BertModel.from_pretrained("bert-base-uncased")
+model = BertForQuestionAnswering.from_pretrained("bert-base-uncased")
 
 def answer_query(question, context):
-    # Preprocess the question and context using the tokenizer
+    # Preprocess using tokenizer
     inputs = tokenizer(question, context, return_tensors="pt")
 
-    # Use the model to get the answer
+    # Use model for question answering
     with torch.no_grad():
         outputs = model(**inputs)
 
-        # Access the logits from the model's output structure
-        start_logits = outputs.hidden_states[-1][:, 0, :]  # Access from hidden states
-        end_logits = outputs.hidden_states[-1][:, 1, :]
+        # Retrieve logits directly
+        start_logits = outputs.start_logits
+        end_logits = outputs.end_logits
 
-    # Find the most likely answer span
+    # Find answer span
     answer_start = torch.argmax(start_logits)
     answer_end = torch.argmax(end_logits) + 1
 
-    # Extract the answer from the context
-    answer = tokenizer.convert_tokens_to_string(context)[answer_start:answer_end]
+    # Extract answer from context
+    answer = tokenizer.convert_tokens_to_string(
+        tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])  # Access original tokens
+    )[answer_start:answer_end]
 
     return answer
 
+
 # Streamlit app
 st.title("Question Answering App")