Update app.py

app.py

@@ -6,10 +6,12 @@ from sklearn.metrics.pairwise import cosine_similarity
 import numpy as np
 import nltk
 from nltk.tokenize import sent_tokenize
+from rank_bm25 import BM25Okapi  # For BM25 retrieval
 
 nltk.download('punkt')
 
-model_name = "deepset/roberta-base-squad2"
+# Use a strong RoBERTa model
+model_name = "deepset/roberta-large-squad2"  # More powerful than base
 model = AutoModelForQuestionAnswering.from_pretrained(model_name)
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 qa_pipeline = pipeline("question-answering", model=model, tokenizer=tokenizer)
@@ -23,13 +25,14 @@ def read_pdf(file):
     except Exception as e:
         return str(e)
 
-def retrieve_relevant_text(question, context, top_n=3):
+def retrieve_relevant_text_bm25(question, sentences, top_n=3):
     try:
-        vectorizer = TfidfVectorizer().fit_transform([question] + context)
-        vectors = vectorizer.toarray()
-        cosine_matrix = cosine_similarity(vectors)
-        similar_ix = np.argsort(cosine_matrix[0])[::-1][1:top_n+1]
-        relevant_texts = [context[ix] for ix in similar_ix]
+        tokenized_corpus = [sent.split() for sent in sentences]
+        bm25 = BM25Okapi(tokenized_corpus)
+        tokenized_query = question.split()
+        doc_scores = bm25.get_scores(tokenized_query)
+        top_n_indices = np.argsort(doc_scores)[::-1][:top_n]  # Get indices of top N
+        relevant_texts = [sentences[i] for i in top_n_indices]
         return " ".join(relevant_texts)
     except Exception as e:
         return str(e)
@@ -41,12 +44,10 @@ def answer_question(pdf, question, num_words):
             return text
 
         sentences = sent_tokenize(text)
-        relevant_text = retrieve_relevant_text(question, sentences)
+        relevant_text = retrieve_relevant_text_bm25(question, sentences) # Use BM25
 
         response = qa_pipeline(question=question, context=relevant_text)
         answer = response['answer']
-        start = response['start']
-        end = response['end']
 
         words = answer.split()
         if len(words) > num_words:
@@ -77,7 +78,7 @@ iface = gr.Interface(
         gr.Slider(minimum=1, maximum=1000, value=100, step=1, label="Number of Words")
     ],
     outputs=gr.Textbox(label="Answer"),
-    title="PDF Q&A | Made by PanigrahiNirma"
+    title="PDF Q&A with RoBERTa | Made by NP"
 )
 
 if __name__ == "__main__":