Upload app.py

app.py

@@ -0,0 +1,50 @@
+import streamlit as st
+import pdfplumber
+from transformers import pipeline, RagTokenizer, RagRetriever, RagSequenceForGeneration
+
+def preprocess_text(text):
+    # Remove extra whitespace and normalize line breaks
+    text = text.replace('\n', ' ').replace('\r', '')
+    text = ' '.join(text.split())
+    return text
+
+st.title("Chat with Your PDF")
+
+uploaded_file = st.file_uploader("Choose a PDF file", type="pdf")
+
+if uploaded_file is not None:
+    with st.spinner('Reading PDF...'):
+        # Extract text from PDF using pdfplumber
+        with pdfplumber.open(uploaded_file) as pdf:
+            text = ""
+            for page in pdf.pages:
+                text += page.extract_text()
+
+        text = preprocess_text(text)
+        st.success('PDF successfully read and preprocessed!')
+
+        # Display some text from the PDF
+        st.text_area("Extracted Text", text[:1000], height=300)
+
+        # Initialize the RAG model
+        tokenizer = RagTokenizer.from_pretrained("facebook/rag-token-nq")
+        retriever = RagRetriever.from_pretrained("facebook/rag-token-nq", use_dummy_dataset=True)
+        rag_model = RagSequenceForGeneration.from_pretrained("facebook/rag-token-nq")
+
+        # Tokenize the text for RAG
+        input_texts = text.split('. ')
+        input_ids = tokenizer(input_texts, return_tensors="pt", padding=True, truncation=True, max_length=512)
+
+        # Build context embeddings for retrieval
+        context_input_ids = retriever(input_ids.input_ids, input_ids.input_ids, num_beams=2)
+
+        question = st.text_input("Ask a question about the PDF:")
+        if question:
+            with st.spinner('Searching for answer...'):
+                # Tokenize the question
+                question_ids = tokenizer(question, return_tensors="pt")['input_ids']
+                
+                # Generate answer using RAG
+                generated = rag_model.generate(input_ids=context_input_ids.input_ids, context_input_ids=question_ids, num_beams=2)
+                rag_answer = tokenizer.decode(generated[0], skip_special_tokens=True)
+                st.write(rag_answer)