Update app.py

app.py

@@ -1,53 +1,63 @@
 import streamlit as st
 import PyPDF2
-from transformers import RagTokenizer, RagRetriever, RagTokenForGeneration
+import os
+from transformers import RagTokenizer, RagRetriever, RagSequenceForGeneration
+import faiss
+import torch
 
-# Load PDF and extract text
+# Function to load PDF and extract text
 def load_pdf(uploaded_file):
     reader = PyPDF2.PdfReader(uploaded_file)
-    text = ""
+    text = ''
     for page in reader.pages:
-        if page.extract_text():  # Ensure text extraction is valid
-            text += page.extract_text() + "\n"
+        text += page.extract_text()
     return text
 
 # Initialize RAG model
 def initialize_rag_model():
-    # Load the tokenizer and model
+    # Load tokenizer, retriever, and model
     tokenizer = RagTokenizer.from_pretrained("facebook/rag-token-nq")
-    
-    # Use a dummy retriever for testing purposes
     retriever = RagRetriever.from_pretrained("facebook/rag-token-nq", use_dummy_dataset=True)
-    
-    model = RagTokenForGeneration.from_pretrained("facebook/rag-token-nq")
+    model = RagSequenceForGeneration.from_pretrained("facebook/rag-token-nq")
     return tokenizer, retriever, model
 
-# Process user query
-def generate_answer(query, context, tokenizer, retriever, model):
-    # Tokenize the input question
-    inputs = tokenizer(query, return_tensors="pt")
+# Function to answer questions
+def answer_question(question, context, tokenizer, model):
+    input_ids = tokenizer.encode(question, return_tensors='pt')
+    context_ids = tokenizer.encode(context, return_tensors='pt')
+    input_ids = input_ids.to(model.device)
+    context_ids = context_ids.to(model.device)
 
-    # Prepare inputs for the model with a dummy context
-    inputs["context_input_ids"] = retriever(context, return_tensors="pt")["input_ids"]
-    
-    # Generate the answer
-    outputs = model.generate(**inputs)
-    answer = tokenizer.batch_decode(outputs, skip_special_tokens=True)
-    return answer[0]
+    # Generate answer
+    with torch.no_grad():
+        outputs = model(input_ids=input_ids, context_input_ids=context_ids)
+    answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return answer
 
-# Streamlit UI
-st.title("PDF Question-Answer Chatbot")
+# Main Streamlit application
+st.title("PDF Q&A Chatbot")
+st.write("Upload a PDF file and ask questions about its content.")
 
-uploaded_file = st.file_uploader("Upload a PDF file", type=["pdf"])
-if uploaded_file is not None:
+# Upload PDF file
+uploaded_file = st.file_uploader("Choose a PDF file", type="pdf")
+
+if uploaded_file:
     text = load_pdf(uploaded_file)
-    st.write("PDF loaded successfully. You can now ask questions.")
-    
-    # Initialize the RAG model
+    st.write("PDF content loaded successfully.")
+
+    # Initialize model
     tokenizer, retriever, model = initialize_rag_model()
+    
+    # Get user question
+    question = st.text_input("Enter your question:")
+
+    if st.button("Get Answer"):
+        if text:
+            # Call the answer_question function
+            answer = answer_question(question, text, tokenizer, model)
+            st.write("Answer:", answer)
+        else:
+            st.error("No text found in the PDF.")
+
 
-    user_query = st.text_input("Ask a question about the PDF:")
-    if user_query:
-        answer = generate_answer(user_query, text, tokenizer, retriever, model)
-        st.write(f"Answer: {answer}")  # Display the answer