Update app.py

app.py

@@ -1,11 +1,13 @@
 import streamlit as st
 import pdfplumber
 import re
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline, DPRQuestionEncoder, DPRContextEncoder
 from gtts import gTTS
 from sklearn.feature_extraction.text import CountVectorizer
 import nltk
 from nltk.sentiment import SentimentIntensityAnalyzer
+import faiss
+import numpy as np
 
 # Download necessary NLTK data
 nltk.download('vader_lexicon')
@@ -16,6 +18,10 @@ model = AutoModelForSeq2SeqLM.from_pretrained("google/pegasus-xsum")
 qa_pipeline = pipeline("question-answering", model="distilbert-base-cased-distilled-squad")
 sia = SentimentIntensityAnalyzer()
 
+# Initialize RAG components
+question_encoder = DPRQuestionEncoder.from_pretrained("facebook/dpr-question_encoder-single-nq-base")
+context_encoder = DPRContextEncoder.from_pretrained("facebook/dpr-ctx_encoder-single-nq-base")
+
 # Helper functions
 def extract_text_from_pdf(file):
     with pdfplumber.open(file) as pdf:
@@ -42,14 +48,30 @@ def summarize_large_document(text, max_length=800):
     summaries = [summarize_text_pegasus(paragraph, max_length=max_length) for paragraph in paragraphs]
     return " ".join(summaries)
 
-def answer_question_with_context(question, context, chunk_size=500):
-    paragraphs = split_text_into_paragraphs(context)
-    answers = []
+def embed_text(text, encoder, tokenizer):
+    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding="longest")
+    embeddings = encoder(**inputs).pooler_output
+    return embeddings.detach().numpy()
+
+def build_index(paragraphs):
+    index = faiss.IndexFlatL2(768)
+    embeddings = []
     for paragraph in paragraphs:
-        try:
-            answers.append(qa_pipeline({'question': question, 'context': paragraph})['answer'])
-        except:
-            continue
+        embeddings.append(embed_text(paragraph, context_encoder, tokenizer))
+    embeddings = np.vstack(embeddings)
+    index.add(embeddings)
+    return index, paragraphs
+
+def retrieve_relevant_paragraphs(question, index, paragraphs, top_k=5):
+    question_embedding = embed_text(question, question_encoder, tokenizer)
+    distances, indices = index.search(question_embedding, top_k)
+    return [paragraphs[i] for i in indices[0]]
+
+def answer_question_with_rag(question, context, top_k=5):
+    paragraphs = split_text_into_paragraphs(context)
+    index, paragraphs = build_index(paragraphs)
+    relevant_paragraphs = retrieve_relevant_paragraphs(question, index, paragraphs, top_k)
+    answers = [qa_pipeline({'question': question, 'context': paragraph})['answer'] for paragraph in relevant_paragraphs]
     return " ".join(answers)
 
 def text_to_speech(text, lang="en"):
@@ -73,6 +95,9 @@ def analyze_sentiment(text):
 st.title("Enhanced PDF to Audiobook App")
 st.markdown("### Turn documents into interactive audiobooks with advanced features.")
 
+if 'audio_path' not in st.session_state:
+    st.session_state['audio_path'] = None
+
 uploaded_file = st.file_uploader("Upload a PDF", type="pdf")
 
 if uploaded_file:
@@ -89,20 +114,24 @@ if uploaded_file:
         if st.button("Convert Summary to Audiobook"):
             with st.spinner("Generating audio..."):
                 audio_path = text_to_speech(summary)
-            st.audio(audio_path, format="audio/mp3")
-            st.download_button("Download Audio", data=open(audio_path, "rb"), file_name="summary_audio.mp3")
+                st.session_state['audio_path'] = audio_path
+        if st.session_state['audio_path']:
+            st.audio(st.session_state['audio_path'], format="audio/mp3")
+            st.download_button("Download Audio", data=open(st.session_state['audio_path'], "rb"), file_name="summary_audio.mp3")
 
     st.markdown("### Ask Questions About the Document")
     question = st.text_input("Your Question:")
     if question:
         with st.spinner("Answering your question..."):
-            answer = answer_question_with_context(question, cleaned_text)
+            answer = answer_question_with_rag(question, cleaned_text)
         st.write(f"**Answer:** {answer}")
         if st.button("Convert Answer to Audio"):
             with st.spinner("Generating answer audio..."):
                 answer_audio_path = text_to_speech(answer)
-            st.audio(answer_audio_path, format="audio/mp3")
-            st.download_button("Download Answer Audio", data=open(answer_audio_path, "rb"), file_name="answer_audio.mp3")
+                st.session_state['audio_path'] = answer_audio_path
+        if st.session_state['audio_path']:
+            st.audio(st.session_state['audio_path'], format="audio/mp3")
+            st.download_button("Download Answer Audio", data=open(st.session_state['audio_path'], "rb"), file_name="answer_audio.mp3")
 
     st.markdown("### Document Insights")
     if st.checkbox("Extract Keywords"):