Update app.py

app.py

@@ -1,5 +1,5 @@
 import streamlit as st
-from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, pipeline, AutoModelForQuestionAnswering
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, AutoModelForQuestionAnswering
 from sentence_transformers import SentenceTransformer
 import fitz  # PyMuPDF
 import os
@@ -9,10 +9,9 @@ summarization_model_name = 'facebook/bart-large-cnn'
 tokenizer = AutoTokenizer.from_pretrained(summarization_model_name)
 summarization_model = AutoModelForSeq2SeqLM.from_pretrained(summarization_model_name)
 
-qa_model_name = 'distilbert-base-uncased-distilled-squad'
+qa_model_name = 'deepset/bert-large-uncased-whole-word-masking-squad2'
 qa_tokenizer = AutoTokenizer.from_pretrained(qa_model_name)
 qa_model = AutoModelForQuestionAnswering.from_pretrained(qa_model_name)
-qa_pipeline = pipeline('question-answering', model=qa_model, tokenizer=qa_tokenizer)
 
 # Function to extract text from a PDF file
 def extract_text_from_pdf(file):
@@ -28,6 +27,15 @@ def summarize_document(document):
     summary_ids = summarization_model.generate(inputs['input_ids'], max_length=150, min_length=30, length_penalty=2.0, num_beams=4, early_stopping=True)
     return tokenizer.decode(summary_ids[0], skip_special_tokens=True)
 
+# Function to get answer to question
+def get_answer(question, context):
+    inputs = qa_tokenizer(question, context, return_tensors="pt")
+    start_positions, end_positions = qa_model(**inputs)
+    answer_start = torch.argmax(start_positions)
+    answer_end = torch.argmax(end_positions) + 1
+    answer = qa_tokenizer.convert_tokens_to_string(qa_tokenizer.convert_ids_to_tokens(inputs["input_ids"][0][answer_start:answer_end]))
+    return answer
+
 # Streamlit app
 st.title("PDF Summarizer and Q&A")
 st.write("Upload a PDF file to get a summary and ask questions about the content.")
@@ -57,9 +65,9 @@ if uploaded_file is not None:
     if st.button("Get Answer"):
         if question:
             with st.spinner('Generating answer...'):
-                answer = qa_pipeline({'question': question, 'context': document_text})
+                answer = get_answer(question, document_text)
                 st.write("**Answer:**")
-                st.write(answer['answer'])
+                st.write(answer)
         else:
             st.write("Please enter a question.")