Update app.py

app.py

@@ -1,24 +1,13 @@
 import streamlit as st
 import torch
 from transformers import RagTokenizer, RagRetriever, RagSequenceForGeneration
-from datasets import load_dataset
 import pandas as pd
 import pdfplumber
 
-# Load RAG model and tokenizer
+# Initialize RAG components
 tokenizer = RagTokenizer.from_pretrained("facebook/rag-sequence-nq")
-retriever = RagRetriever.from_pretrained("facebook/rag-sequence-nq", index_name="exact")
-model = RagSequenceForGeneration.from_pretrained("facebook/rag-sequence-nq", retriever=retriever)
-
-# Load the wiki_dpr dataset with trust_remote_code=True
-dataset = load_dataset("facebook/wiki_dpr", split="train", trust_remote_code=True)
-
-# Function to get RAG embeddings
-def get_rag_embeddings(question, context):
-    inputs = tokenizer(question, context, return_tensors="pt", truncation=True)
-    with torch.no_grad():
-        output = model.generate(input_ids=inputs['input_ids'], attention_mask=inputs['attention_mask'])
-    return tokenizer.batch_decode(output, skip_special_tokens=True)[0]
+retriever = RagRetriever.from_pretrained("facebook/rag-sequence-nq")
+model = RagSequenceForGeneration.from_pretrained("facebook/rag-sequence-nq")
 
 # Extract text from PDF
 def extract_text_from_pdf(pdf_file):
@@ -26,19 +15,16 @@ def extract_text_from_pdf(pdf_file):
         text = ""
         for page in pdf.pages:
             page_text = page.extract_text()
-            if page_text:  # Check if the page has extractable text
+            if page_text:
                 text += page_text + "\n"
-    return text.strip()  # Return stripped text for better formatting
+    return text.strip()
 
-# Store the PDF text and CSV data
-pdf_text = ""
-csv_data = None
-
-# Streamlit app UI
+# Streamlit UI
 st.title("RAG-Powered PDF & CSV Chatbot")
 
 # CSV file upload
 csv_file = st.file_uploader("Upload a CSV file", type=["csv"])
+csv_data = None
 if csv_file:
     csv_data = pd.read_csv(csv_file)
     st.write("CSV file loaded successfully!")
@@ -46,6 +32,7 @@ if csv_file:
 
 # PDF file upload
 pdf_file = st.file_uploader("Upload a PDF file", type=["pdf"])
+pdf_text = ""
 if pdf_file:
     pdf_text = extract_text_from_pdf(pdf_file)
     if pdf_text:
@@ -62,15 +49,14 @@ if st.button("Get Response"):
     if not pdf_text and csv_data is None:
         st.warning("Please upload a PDF or CSV file first.")
     else:
-        # Combine PDF text and CSV content for context in RAG
         combined_context = pdf_text
         if csv_data is not None:
             combined_context += "\n" + csv_data.to_string()
 
-        # Get RAG-generated response
-        try:
-            response = get_rag_embeddings(user_input, combined_context)
-            st.write("### Response:")
-            st.write(response)
-        except Exception as e:
-            st.error(f"Error while processing the question: {e}")
+        # Generate response using RAG
+        inputs = tokenizer(user_input, combined_context, return_tensors="pt", truncation=True)
+        with torch.no_grad():
+            output = model.generate(input_ids=inputs['input_ids'], attention_mask=inputs['attention_mask'])
+        response = tokenizer.batch_decode(output, skip_special_tokens=True)[0]
+        st.write("### Response:")
+        st.write(response)