Update app.py

app.py

@@ -6,7 +6,7 @@ from PyPDF2 import PdfReader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.embeddings import HuggingFaceEmbeddings
 from langchain.vectorstores import FAISS
-from transformers import pipeline, AutoModel, AutoTokenizer
+from transformers import pipeline
 import torch
 
 # Set up the page configuration as the first Streamlit command
@@ -15,16 +15,29 @@ 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")  # Use a summarization model
+    summarizer = pipeline("summarization", model="facebook/bart-large-cnn")
     return summarizer
 
 summarizer = load_summarization_pipeline()
 
-# List of Hugging Face PDF URLs
-PDF_URLS = [
-    "https://huggingface.co/spaces/tahirsher/GenAI_Lawyers_Guide/blob/main/administrator92ada0936848e501425591b4ad0cd417.pdf",
-    "https://huggingface.co/spaces/tahirsher/GenAI_Lawyers_Guide/blob/main/Pakistan%20Penal%20Code.pdf",
-]
+# Dictionary of Hugging Face PDF URLs grouped by folders
+PDF_FOLDERS = {
+    "PPC and Administration": [
+        "https://huggingface.co/spaces/tahirsher/GenAI_Lawyers_Guide/blob/main/administrator92ada0936848e501425591b4ad0cd417.pdf",
+        "https://huggingface.co/spaces/tahirsher/GenAI_Lawyers_Guide/blob/main/Pakistan%20Penal%20Code.pdf",
+    ],
+    "IHC": [
+        "https://huggingface.co/spaces/tahirsher/GenAI_Lawyers_Guide/tree/main/IHC"
+    "LHC": [
+        "https://huggingface.co/spaces/tahirsher/GenAI_Lawyers_Guide/tree/main/LHC"
+    "Lahore High Court Rules and Orders": [
+        "https://huggingface.co/spaces/tahirsher/GenAI_Lawyers_Guide/tree/main/Lahore%20High%20Court%20Rules%20and%20Orders"
+    "PHC": [
+        "https://huggingface.co/spaces/tahirsher/GenAI_Lawyers_Guide/tree/main/PHC"
+    "SC": [
+        "https://huggingface.co/spaces/tahirsher/GenAI_Lawyers_Guide/tree/main/SC"
+    ],
+}
 
 # Helper function to convert Hugging Face blob URLs to direct download URLs
 def get_huggingface_raw_url(url):
@@ -32,25 +45,28 @@ def get_huggingface_raw_url(url):
         return url.replace("/blob/", "/resolve/")
     return url
 
-# Fetch and extract text from PDF files hosted on Hugging Face
-def fetch_pdf_text_from_huggingface(urls):
-    text = ""
-    for url in urls:
-        raw_url = get_huggingface_raw_url(url)
-        response = requests.get(raw_url)
-        if response.status_code == 200:
-            pdf_file = BytesIO(response.content)
-            try:
-                pdf_reader = PdfReader(pdf_file)
-                for page in pdf_reader.pages:
-                    page_text = page.extract_text()
-                    if page_text:
-                        text += page_text
-            except Exception as e:
-                st.error(f"Failed to read PDF from URL {url}: {e}")
-        else:
-            st.error(f"Failed to fetch PDF from URL: {url}")
-    return text
+# Fetch and extract text from all PDFs in specified folders
+def fetch_pdf_text_from_folders(pdf_folders):
+    all_text = ""
+    for folder_name, urls in pdf_folders.items():
+        folder_text = f"\n[Folder: {folder_name}]\n"
+        for url in urls:
+            raw_url = get_huggingface_raw_url(url)
+            response = requests.get(raw_url)
+            if response.status_code == 200:
+                pdf_file = BytesIO(response.content)
+                try:
+                    pdf_reader = PdfReader(pdf_file)
+                    for page in pdf_reader.pages:
+                        page_text = page.extract_text()
+                        if page_text:
+                            folder_text += page_text
+                except Exception as e:
+                    st.error(f"Failed to read PDF from URL {url}: {e}")
+            else:
+                st.error(f"Failed to fetch PDF from URL: {url}")
+        all_text += folder_text
+    return all_text
 
 # Split text into manageable chunks
 @st.cache_data
@@ -70,14 +86,9 @@ def load_or_create_vector_store(text_chunks):
 
 # Generate summary based on the retrieved text
 def generate_summary_with_huggingface(query, retrieved_text):
-    # Concatenate query and retrieved text for summarization
     summarization_input = f"{query}\n\nRelated information:\n{retrieved_text}"
-    
-    # Truncate input to fit within the model’s token length limit (approximately 1024 tokens)
     max_input_length = 1024
     summarization_input = summarization_input[:max_input_length]
-
-    # Generate the summary
     summary = summarizer(summarization_input, max_length=500, min_length=50, do_sample=False)
     return summary[0]["summary_text"]
 
@@ -90,13 +101,10 @@ def user_input(user_question, vector_store):
 # Main function to run the Streamlit app
 def main():
     st.title("📄 Gen AI Lawyers Guide")
-
-    # Load documents from Hugging Face
-    raw_text = fetch_pdf_text_from_huggingface(PDF_URLS)
+    raw_text = fetch_pdf_text_from_folders(PDF_FOLDERS)
     text_chunks = get_text_chunks(raw_text)
     vector_store = load_or_create_vector_store(text_chunks)
 
-    # User question input
     user_question = st.text_input("Ask a Question:", placeholder="Type your question here...")
 
     if st.button("Get Response"):
@@ -108,4 +116,4 @@ def main():
                 st.markdown(f"**🤖 AI:** {answer}")
 
 if __name__ == "__main__":
-    main()
+    main()