Update app.py

app.py

@@ -1,26 +1,93 @@
-if st.button("Get Answer"):
+import streamlit as st
+from transformers import DPRQuestionEncoder, DPRQuestionEncoderTokenizer
+from transformers import BartForConditionalGeneration, BartTokenizer
+from sentence_transformers import SentenceTransformer
+import pdfplumber
+import numpy as np
+from sklearn.metrics.pairwise import cosine_similarity
+import torch
+import gc
+
+# Load the Question Encoder, Context Encoder, and Tokenizers
+question_encoder = DPRQuestionEncoder.from_pretrained("facebook/dpr-question_encoder-single-nq-base")
+question_tokenizer = DPRQuestionEncoderTokenizer.from_pretrained("facebook/dpr-question_encoder-single-nq-base")
+
+# Load the Generator Model
+generator = BartForConditionalGeneration.from_pretrained("facebook/bart-large-cnn")
+generator_tokenizer = BartTokenizer.from_pretrained("facebook/bart-large-cnn")
+
+# Load Sentence Embedding Model for Vector Store
+sentence_model = SentenceTransformer('paraphrase-MiniLM-L6-v2')
+
+# Initialize documents list with some sample documents
+documents = [
+    "Streamlit is an open-source Python library that makes it easy to build beautiful custom web-apps for machine learning and data science.",
+    "Hugging Face is a company that provides tools and models for natural language processing (NLP).",
+    "Retrieval-Augmented Generation (RAG) is a method that combines document retrieval with a generative model for question answering.",
+]
+
+# Encode the initial documents for similarity comparison
+doc_embeddings = sentence_model.encode(documents)
+
+# Streamlit Frontend
+st.set_page_config(page_title="RAG-based PDF Query Application", layout="wide")
+st.title("📄 Retrieval-Augmented Generation (RAG) Application")
+
+# File Upload for PDF Documents
+uploaded_file = st.file_uploader("Upload a PDF file", type="pdf")
+if uploaded_file:
+    # Extract text from PDF
+    pdf_text = ""
+    with pdfplumber.open(uploaded_file) as pdf:
+        for page_num, page in enumerate(pdf.pages):
+            if page_num > 20:  # Limit to first 20 pages for efficiency
+                break
+            page_text = page.extract_text()
+            if page_text:  # Check if text was extracted
+                pdf_text += page_text + " "
+
+    if pdf_text:
+        # Add the PDF text to the documents list and update document embeddings
+        documents.append(pdf_text)
+        pdf_embedding = sentence_model.encode([pdf_text])
+        doc_embeddings = np.vstack([doc_embeddings, pdf_embedding])
+        st.success("PDF text added to knowledge base for querying!")
+    else:
+        st.error("No text could be extracted from the PDF.")
+
+# User Input
+st.markdown("Enter your query below:")
+query = st.text_input("🔍 Enter your query")
+
+if st.button("💬 Get Answer"):
     if query:
         try:
-            # Existing embedding process
+            # Step 1: Encode the query
             question_inputs = question_tokenizer(query, return_tensors="pt")
             question_embedding = question_encoder(**question_inputs).pooler_output.detach().cpu().numpy()
 
-            # Cosine similarity
+            # Step 2: Calculate Cosine Similarity
             similarity_scores = cosine_similarity(question_embedding, doc_embeddings)
+
+            # Step 3: Get the indices of the top 3 most similar documents
             top_indices = similarity_scores[0].argsort()[-3:][::-1]
             retrieved_docs = [documents[idx] for idx in top_indices]
+
+            # Step 4: Concatenate retrieved documents
             context = " ".join(retrieved_docs)
 
             # Log the retrieved context for debugging
             st.write(f"Context for the query: {context}")
 
-            # Ensure the question and context are correctly formatted
+            # Step 5: Use the Generator to Answer the Question
             input_ids = generator_tokenizer.encode(f"question: {query} context: {context}", return_tensors="pt")
-            outputs = generator.generate(input_ids, max_length=200)
+            outputs = generator.generate(input_ids, max_length=200, num_return_sequences=1)
+
+            # Decode and display the response
             answer = generator_tokenizer.decode(outputs[0], skip_special_tokens=True)
-            
             st.write("**Answer:**")
             st.write(answer)
+
         except Exception as e:
             st.error(f"An error occurred: {str(e)}")
         finally: