app.py

import os
import streamlit as st
from io import BytesIO
from PyPDF2 import PdfReader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import FAISS
from transformers import pipeline
import torch

# Set up the page configuration
st.set_page_config(page_title="RAG-based PDF Chat", layout="centered", page_icon="📄")

# Load the summarization pipeline model
@st.cache_resource
def load_summarization_pipeline():
    summarizer = pipeline("summarization", model="facebook/bart-large-cnn")
    return summarizer

summarizer = load_summarization_pipeline()

# Helper function to extract text from PDFs in a local folder
def fetch_pdf_text_from_folder(folder_path):
    all_text = ""
    pdf_files = [f for f in os.listdir(folder_path) if f.endswith('.pdf')]
    total_files = len(pdf_files)
    
    if total_files == 0:
        st.warning("No PDF files found in the folder.")
        return ""

    progress_bar = st.progress(0)
    for index, file_name in enumerate(pdf_files):
        try:
            file_path = os.path.join(folder_path, file_name)
            with open(file_path, 'rb') as file:
                pdf_reader = PdfReader(file)
                for page in pdf_reader.pages:
                    page_text = page.extract_text()
                    if page_text:
                        all_text += f"\n[File: {file_name}]\n{page_text}"
        except Exception as e:
            st.error(f"Failed to read PDF file {file_name}: {e}")
        
        # Update the progress bar
        progress_percentage = int(((index + 1) / total_files) * 100)
        progress_bar.progress(progress_percentage)
    
    return all_text

# Split text into manageable chunks
@st.cache_data
def get_text_chunks(text):
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=10000, chunk_overlap=1000)
    chunks = text_splitter.split_text(text)
    return chunks

# Initialize embedding function
embedding_function = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")

# Create a FAISS vector store with embeddings, checking for empty chunks
@st.cache_resource
def load_or_create_vector_store(text_chunks):
    if not text_chunks:
        st.error("No valid text chunks found to create a vector store. Please check your PDF files or content.")
        return None
    vector_store = FAISS.from_texts(text_chunks, embedding=embedding_function)
    return vector_store

# Generate summary based on the retrieved text
def generate_summary_with_huggingface(query, retrieved_text):
    summarization_input = f"{query}\n\nRelated information:\n{retrieved_text}"
    max_input_length = 1024
    summarization_input = summarization_input[:max_input_length]
    summary = summarizer(summarization_input, max_length=500, min_length=50, do_sample=False)
    return summary[0]["summary_text"]

# Generate response for user query
def user_input(user_question, vector_store):
    if vector_store is None:
        return "Vector store is empty due to failed PDF loading or empty documents."
    docs = vector_store.similarity_search(user_question)
    context_text = " ".join([doc.page_content for doc in docs])
    return generate_summary_with_huggingface(user_question, context_text)

# Main function to run the Streamlit app
def main():
    st.title("📄 Gen AI Lawyers Guide")
    st.info("Loading data from the 'law-docs' folder...")

    folder_path = "law-docs"
    raw_text = fetch_pdf_text_from_folder(folder_path)
    text_chunks = get_text_chunks(raw_text)
    vector_store = load_or_create_vector_store(text_chunks)

    user_question = st.text_input("Ask a Question:", placeholder="Type your question here...")

    if st.button("Get Response"):
        if not user_question:
            st.warning("Please enter a question before submitting.")
        else:
            with st.spinner("Generating response..."):
                answer = user_input(user_question, vector_store)
                st.markdown(f"**🤖 AI:** {answer}")

if __name__ == "__main__":
    main()