greyskyAI/PDF_RAG_with_Citations

1import streamlit as st2import pdfplumber3from langchain_community.embeddings import HuggingFaceEmbeddings4from langchain_community.vectorstores import FAISS5from langchain.text_splitter import RecursiveCharacterTextSplitter6import os7from langchain.schema import Document8 9# Set page config10st.set_page_config(page_title="RAG PDF App", layout="wide")11 12# Title13st.title("📄 Retrieval-Augmented Generation (RAG) PDF App")14 15# Sidebar for Parameters16st.sidebar.header("Processing Parameters")17 18chunk_size = st.sidebar.number_input(19    "Chunk Size",20    min_value=100,21    max_value=2000,22    value=500,23    step=100,24    help="The size of each text chunk."25)26 27chunk_overlap = st.sidebar.number_input(28    "Chunk Overlap",29    min_value=0,30    max_value=500,31    value=50,32    step=10,33    help="The number of characters to overlap between chunks."34)35 36use_rag = st.sidebar.checkbox(37    "Use Retrieval-Augmented Generation (RAG)",38    value=True,39    help="Enable or disable RAG functionality."40)41 42num_chunks_retrieved = st.sidebar.number_input(43    "Number of Chunks to Retrieve",44    min_value=1,45    max_value=20,46    value=5,47    step=1,48    help="The number of relevant chunks to retrieve for each query."49)50 51# Initialize embeddings52embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")53 54def process_uploaded_files(uploaded_files):55    """56    Processes uploaded PDF files by extracting text and creating Document objects with metadata.57    58    Args:59        uploaded_files (list): List of uploaded PDF files.60    61    Returns:62        list: List of Document objects containing page content and metadata.63        list: List of filenames corresponding to the uploaded files.64    """65    documents = []66    filenames = []67    68    for uploaded_file in uploaded_files:69        # Save the uploaded file to the uploads directory70        file_path = os.path.join("uploads", uploaded_file.name)71        with open(file_path, "wb") as f:72            f.write(uploaded_file.getbuffer())73 74        # Extract text using pdfplumber75        try:76            with pdfplumber.open(file_path) as pdf:77                for page_number, page in enumerate(pdf.pages, start=1):78                    page_text = page.extract_text()79                    if page_text:80                        # Create Document object with metadata81                        document = Document(82                            page_content=page_text,83                            metadata={84                                "source": uploaded_file.name,85                                "page_number": page_number86                            }87                        )88                        documents.append(document)89        except Exception as e:90            st.error(f"Error processing {uploaded_file.name}: {e}")91            continue92 93        # Append to filenames list94        filenames.append(uploaded_file.name)95    96    return documents, filenames97 98# File Uploader99uploaded_files = st.file_uploader("Upload PDF Files", type=["pdf"], accept_multiple_files=True)100 101# Button to Process PDFs102if st.button("Process PDFs"):103    if not uploaded_files:104        st.warning("Please upload at least one PDF file.")105    else:106        # Process uploaded files107        documents, filenames = process_uploaded_files(uploaded_files)108 109        # Display Uploaded Files110        st.subheader("Uploaded Files")111        for filename in filenames:112            st.write(f"- {filename}")113 114        # Process Texts with or without RAG115        if use_rag:116            st.subheader("Processing with RAG")117            with st.spinner("Creating embeddings and setting up vector store..."):118                splitter = RecursiveCharacterTextSplitter(119                    chunk_size=chunk_size,120                    chunk_overlap=chunk_overlap121                )122                123                # Split documents124                chunks = splitter.split_documents(documents)125 126                # Create FAISS vector store127                vector_store = FAISS.from_documents(chunks, embeddings)128 129                # Save the vector store locally130                vector_store.save_local("faiss_index")131 132            st.success("RAG setup complete!")133        else:134            st.subheader("Processing without RAG")135            st.write("All PDFs have been processed. You can now query the entire content.")136 137# Query Interface138st.header("❓ Ask a Question")139 140query = st.text_input("Enter your question here:")141 142if st.button("Submit Query"):143    if not query:144        st.warning("Please enter a question.")145    else:146        if use_rag:147            # Load FAISS vector store148            if not os.path.exists("faiss_index"):149                st.error("No vector store found. Please upload and process PDFs first.")150            else:151                with st.spinner("Retrieving relevant information..."):152                    vector_store = FAISS.load_local(153                        "faiss_index", 154                        embeddings, 155                        allow_dangerous_deserialization=True  # Add this parameter156                    )157                    docs = vector_store.similarity_search(query, k=num_chunks_retrieved)158        159                st.subheader("🔍 Retrieved Chunks")160                for idx, doc in enumerate(docs, start=1):161                    st.markdown(f"**Chunk {idx}:**")162                    st.markdown(f"**Source:** {doc.metadata['source']} | **Page:** {doc.metadata.get('page_number', 'N/A')}")163                    with st.expander("View Chunk"):164                        st.write(doc.page_content)165        else:166            st.write("RAG is disabled. Please enable RAG to retrieve relevant chunks.")167 168# Cleanup: Remove temporary files169for filename in os.listdir("uploads"):170    file_path = os.path.join("uploads", filename)171    os.remove(file_path)