Update app.py

app.py

@@ -1,28 +1,15 @@
 import os
 import streamlit as st
-from transformers import BertTokenizer, BertForQuestionAnswering, pipeline
+from transformers import pipeline
 from PyPDF2 import PdfReader
 import tempfile
 
 # Function to perform question-answering
 @st.cache_data(show_spinner=False)
-def question_answering_bert(questions, pdf_text):
-    tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
-    model = BertForQuestionAnswering.from_pretrained('bert-base-uncased')
-
-    answers = []
-
-    for question in questions:
-        inputs = tokenizer(question, pdf_text, padding=True, return_tensors='pt')
-        outputs = model(**inputs)
-        start_scores = outputs.start_logits
-        end_scores = outputs.end_logits
-
-        start_index = start_scores.argmax()
-        end_index = end_scores.argmax() + 1
-
-        answer = tokenizer.decode(inputs['input_ids'][0][start_index:end_index])
-        answers.append({"answer": answer, "score": start_scores.max().item() + end_scores.max().item()})
+def question_answering(questions, pdf_text):
+    # Perform question-answering using Hugging Face's Transformers
+    question_answerer = pipeline("question-answering", model="distilbert-base-cased-distilled-squad", tokenizer="distilbert-base-cased-distilled-squad")
+    answers = question_answerer(question=questions, context=pdf_text)
 
     return answers
 
@@ -43,8 +30,8 @@ def main():
         pdf_reader = PdfReader(pdf_path)
         pdf_text = "\n".join([pdf_page.extract_text() for pdf_page in pdf_reader.pages])
 
-        # Perform question-answering using BERT model
-        answers = question_answering_bert(questions, pdf_text)
+        # Perform question-answering in batches
+        answers = question_answering(questions, pdf_text)
 
         st.write("Questions and Answers:")
         for i, (question, answer) in enumerate(zip(questions, answers)):