Update app.py

app.py

@@ -4,23 +4,27 @@ from pdfminer.high_level import extract_text
 from sklearn.feature_extraction.text import TfidfVectorizer
 from sklearn.metrics.pairwise import cosine_similarity
 import numpy as np
+import nltk
+from nltk.tokenize import sent_tokenize
 
-# Load BERT for QA
-bert_model_name = "bert-large-uncased-whole-word-masking-finetuned-squad"
-bert_model = AutoModelForQuestionAnswering.from_pretrained(bert_model_name)
-bert_tokenizer = AutoTokenizer.from_pretrained(bert_model_name)
-qa_pipeline = pipeline("question-answering", model=bert_model, tokenizer=bert_tokenizer)
+nltk.download('punkt')  # Download necessary NLTK data
+
+# Use a potentially better QA model
+model_name = "deepset/roberta-base-squad2"  # More robust than the previous one
+model = AutoModelForQuestionAnswering.from_pretrained(model_name)
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+qa_pipeline = pipeline("question-answering", model=model, tokenizer=tokenizer)
 
 def read_pdf(file):
     try:
         text = extract_text(file)
         if not text:
-            raise ValueError("PDF extraction failed. The PDF might be scanned or have an unsupported format.")
+            raise ValueError("PDF extraction failed.")
         return text
     except Exception as e:
         return str(e)
 
-def retrieve_relevant_text(question, context, top_n=5):
+def retrieve_relevant_text(question, context, top_n=3): #reduced the top n
     try:
         vectorizer = TfidfVectorizer().fit_transform([question] + context)
         vectors = vectorizer.toarray()
@@ -33,12 +37,12 @@ def retrieve_relevant_text(question, context, top_n=5):
 
 def answer_question(pdf, question, num_words):
     try:
-        context = read_pdf(pdf).split("\n")
-        if isinstance(context, str):
-            return context  # Return error message if read_pdf failed
-        relevant_text = retrieve_relevant_text(question, context)
-        if isinstance(relevant_text, str):
-            return relevant_text  # Return error message if retrieve_relevant_text failed
+        text = read_pdf(pdf)
+        if isinstance(text, str):
+            return text
+        
+        sentences = sent_tokenize(text) #tokenize the text into sentences
+        relevant_text = retrieve_relevant_text(question, sentences)
 
         response = qa_pipeline(question=question, context=relevant_text)
         answer = response['answer']
@@ -48,21 +52,21 @@ def answer_question(pdf, question, num_words):
             answer = " ".join(words[:num_words])
         elif len(words) < num_words:
             remaining_words = num_words - len(words)
-            sentences = relevant_text.split(". ")
-            added_words = 0
-            for sentence in sentences:
-                sentence_words = sentence.split()
-                if added_words < remaining_words:
-                    words.extend(sentence_words)
-                    added_words += len(sentence_words)
-            answer = " ".join(words[:num_words])
-        return answer
-    except IndexError as e:
-        return f"Error: Could not retrieve enough context. {e}" #Handle IndexErrors
+            added_sentences = []
+            for sentence in sent_tokenize(relevant_text):
+                if remaining_words > 0:
+                    sentence_words = sentence.split()
+                    to_add = min(remaining_words, len(sentence_words))
+                    added_sentences.append(" ".join(sentence_words[:to_add]))
+                    remaining_words -= to_add
+                else:
+                    break
+            answer += " " + " ".join(added_sentences)
+
+        return answer.strip() #strip white spaces
     except Exception as e:
         return str(e)
 
-# Define Gradio interface
 iface = gr.Interface(
     fn=answer_question,
     inputs=[