Update app.py

app.py

@@ -7,22 +7,25 @@ tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
 model = BertModel.from_pretrained("bert-base-uncased")
 
 def answer_query(question, context):
-  # Preprocess the question and context using the tokenizer
-  inputs = tokenizer(question, context, return_tensors="pt")
+    # Preprocess the question and context using the tokenizer
+    inputs = tokenizer(question, context, return_tensors="pt")
 
-  # Use the model to get the answer
-  with torch.no_grad():
-    outputs = model(**inputs)
-    start_scores, end_scores = outputs.start_logits, outputs.end_scores
+    # Use the model to get the answer
+    with torch.no_grad():
+        outputs = model(**inputs)
 
-  # Find the most likely answer span
-  answer_start = torch.argmax(start_scores)
-  answer_end = torch.argmax(end_scores) + 1
+        # 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, :]
 
-  # Extract the answer from the context
-  answer = tokenizer.convert_tokens_to_string(context)[answer_start:answer_end]
+    # Find the most likely answer span
+    answer_start = torch.argmax(start_logits)
+    answer_end = torch.argmax(end_logits) + 1
 
-  return answer
+    # Extract the answer from the context
+    answer = tokenizer.convert_tokens_to_string(context)[answer_start:answer_end]
+
+    return answer
 
 # Streamlit app
 st.title("Question Answering App")
@@ -34,15 +37,15 @@ user_query = st.text_input("Enter your question:")
 uploaded_file = st.file_uploader("Upload a context file (txt):")
 
 if uploaded_file is not None:
-  # Read the uploaded file content
-  context = uploaded_file.read().decode("utf-8")
+    # Read the uploaded file content
+    context = uploaded_file.read().decode("utf-8")
 else:
-  # Use default context if no file uploaded
-  context = "This is a sample context for demonstration purposes. You can upload your own text file for context."
+    # Use default context if no file uploaded
+    context = "This is a sample context for demonstration purposes. You can upload your own text file for context."
 
 # Answer the query if a question is provided
 if user_query:
-  answer = answer_query(user_query, context)
-  st.write(f"Answer: {answer}")
+    answer = answer_query(user_query, context)
+    st.write(f"Answer: {answer}")
 else:
-  st.write("Please enter a question.")
+    st.write("Please enter a question.")