Update app.py

app.py

@@ -1,10 +1,15 @@
 import gradio as gr
 import os
+from pathlib import Path
+import json
+import csv
+import pandas as pd
+from tqdm import tqdm
 api_token = os.getenv("HF_TOKEN")
 
 
 from langchain_community.vectorstores import FAISS
-from langchain_community.document_loaders import PyPDFLoader
+from langchain_community.document_loaders import PyPDFLoader, TextLoader, CSVLoader, JSONLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain_community.vectorstores import Chroma
 from langchain.chains import ConversationalRetrievalChain
@@ -20,56 +25,130 @@ import torch
 list_llm = ["meta-llama/Meta-Llama-3-8B-Instruct", "mistralai/Mistral-7B-Instruct-v0.2"]  
 list_llm_simple = [os.path.basename(llm) for llm in list_llm]
 
-# Load and split PDF document
-def load_doc(list_file_path):
-    # Processing for one document only
-    # loader = PyPDFLoader(file_path)
-    # pages = loader.load()
-    loaders = [PyPDFLoader(x) for x in list_file_path]
-    pages = []
-    for loader in loaders:
-        pages.extend(loader.load())
-    text_splitter = RecursiveCharacterTextSplitter(
-        chunk_size = 1024, 
-        chunk_overlap = 64 
-    )  
-    doc_splits = text_splitter.split_documents(pages)
+# Load and split documents of various types
+def load_doc(list_file_path, progress=gr.Progress()):
+    doc_splits = []
+    
+    progress(0, desc="Preparing to load documents")
+    total_files = len(list_file_path)
+    
+    for i, file_path in enumerate(list_file_path):
+        progress((i/total_files) * 0.5, desc=f"Loading {Path(file_path).name}")
+        file_ext = Path(file_path).suffix.lower()
+        
+        try:
+            # PDF documents
+            if file_ext == '.pdf':
+                loader = PyPDFLoader(file_path)
+                pages = loader.load()
+                doc_splits.extend(split_documents(pages))
+            
+            # Text-based documents
+            elif file_ext in ['.txt', '.md', '.py', '.js', '.html', '.css']:
+                loader = TextLoader(file_path)
+                documents = loader.load()
+                doc_splits.extend(split_documents(documents))
+            
+            # CSV files
+            elif file_ext == '.csv':
+                loader = CSVLoader(file_path)
+                documents = loader.load()
+                doc_splits.extend(split_documents(documents))
+            
+            # JSON files
+            elif file_ext in ['.json', '.jsonl']:
+                # For JSON, we need to determine if it's JSON or JSONL
+                with open(file_path, 'r') as f:
+                    content = f.read().strip()
+                    if content.startswith('[') or content.startswith('{'):
+                        # Regular JSON
+                        loader = JSONLoader(
+                            file_path=file_path,
+                            jq_schema='.',
+                            text_content=False
+                        )
+                        documents = loader.load()
+                        doc_splits.extend(split_documents(documents))
+                    else:
+                        # JSONL - process line by line
+                        documents = []
+                        with open(file_path, 'r') as f:
+                            for line in f:
+                                if line.strip():
+                                    try:
+                                        json_obj = json.loads(line)
+                                        text = json.dumps(json_obj)
+                                        documents.append(text)
+                                    except json.JSONDecodeError:
+                                        continue
+                        
+                        text_splitter = RecursiveCharacterTextSplitter(
+                            chunk_size=1024,
+                            chunk_overlap=64
+                        )
+                        doc_splits.extend(text_splitter.create_documents(documents))
+        except Exception as e:
+            print(f"Error processing {file_path}: {str(e)}")
+            continue
+            
+        progress(0.5 + (i/total_files) * 0.5, desc=f"Processed {Path(file_path).name}")
+            
     return doc_splits
 
+# Helper function to split documents
+def split_documents(documents):
+    text_splitter = RecursiveCharacterTextSplitter(
+        chunk_size=1024, 
+        chunk_overlap=64
+    )
+    return text_splitter.split_documents(documents)
+
 # Create vector database
-def create_db(splits):
+def create_db(splits, progress=gr.Progress()):
+    progress(0, desc="Creating vector database")
     embeddings = HuggingFaceEmbeddings()
-    vectordb = FAISS.from_documents(splits, embeddings)
+    
+    # Create vectors with progress bar
+    total_chunks = len(splits)
+    vectordb = FAISS.from_documents(
+        documents=splits,
+        embedding=embeddings
+    )
+    
+    progress(1.0, desc="Vector database creation complete")
     return vectordb
 
 
 # Initialize langchain LLM chain
-# @spaces.GPU(duration=60)
 def initialize_llmchain(llm_model, temperature, max_tokens, top_k, vector_db, progress=gr.Progress()):
+    progress(0, desc=f"Initializing {llm_model}")
+    
     if llm_model == "meta-llama/Meta-Llama-3-8B-Instruct":
         llm = HuggingFaceEndpoint(
             repo_id=llm_model,
-            huggingfacehub_api_token = api_token,
-            temperature = temperature,
-            max_new_tokens = max_tokens,
-            top_k = top_k,
+            huggingfacehub_api_token=api_token,
+            temperature=temperature,
+            max_new_tokens=max_tokens,
+            top_k=top_k,
         )
     else:
         llm = HuggingFaceEndpoint(
-            huggingfacehub_api_token = api_token,
+            huggingfacehub_api_token=api_token,
             repo_id=llm_model, 
-            temperature = temperature,
-            max_new_tokens = max_tokens,
-            top_k = top_k,
+            temperature=temperature,
+            max_new_tokens=max_tokens,
+            top_k=top_k,
         )
     
+    progress(0.5, desc="Setting up memory and retriever")
+    
     memory = ConversationBufferMemory(
         memory_key="chat_history",
         output_key='answer',
         return_messages=True
     )
 
-    retriever=vector_db.as_retriever()
+    retriever = vector_db.as_retriever()
     qa_chain = ConversationalRetrievalChain.from_llm(
         llm,
         retriever=retriever,
@@ -78,27 +157,46 @@ def initialize_llmchain(llm_model, temperature, max_tokens, top_k, vector_db, pr
         return_source_documents=True,
         verbose=False,
     )
+    
+    progress(1.0, desc="LLM chain initialized")
     return qa_chain
 
 # Initialize database
-# @spaces.GPU(duration=60)
 def initialize_database(list_file_obj, progress=gr.Progress()):
     # Create a list of documents (when valid)
     list_file_path = [x.name for x in list_file_obj if x is not None]
+    
+    if not list_file_path:
+        return None, "No valid files uploaded. Please upload at least one file."
+    
     # Load document and create splits
-    doc_splits = load_doc(list_file_path)
+    doc_splits = load_doc(list_file_path, progress)
+    
+    if not doc_splits:
+        return None, "Could not extract any text from the uploaded files."
+    
     # Create or load vector database
-    vector_db = create_db(doc_splits)
-    return vector_db, "Database created!"
+    vector_db = create_db(doc_splits, progress)
+    
+    # Count documents by type
+    file_types = {}
+    for path in list_file_path:
+        ext = Path(path).suffix.lower()
+        file_types[ext] = file_types.get(ext, 0) + 1
+    
+    file_type_summary = ", ".join([f"{count} {ext}" for ext, count in file_types.items()])
+    
+    return vector_db, f"Database created with {len(doc_splits)} chunks from {len(list_file_path)} files ({file_type_summary})!"
 
 # Initialize LLM
-# @spaces.GPU(duration=60)
 def initialize_LLM(llm_option, llm_temperature, max_tokens, top_k, vector_db, progress=gr.Progress()):
-    # print("llm_option",llm_option)
+    if vector_db is None:
+        return None, "Please create a vector database first!"
+    
     llm_name = list_llm[llm_option]
-    print("llm_name: ",llm_name)
+    print("llm_name: ", llm_name)
     qa_chain = initialize_llmchain(llm_name, llm_temperature, max_tokens, top_k, vector_db, progress)
-    return qa_chain, "QA chain initialized. Chatbot is ready!"
+    return qa_chain, f"QA chain initialized with {llm_name}. Chatbot is ready!"
 
 
 def format_chat_history(message, chat_history):
@@ -108,73 +206,166 @@ def format_chat_history(message, chat_history):
         formatted_chat_history.append(f"Assistant: {bot_message}")
     return formatted_chat_history
     
-# @spaces.GPU(duration=60)
 def conversation(qa_chain, message, history):
+    if qa_chain is None:
+        return None, gr.update(value=""), history, "Please initialize the chatbot first!", 0, "", 0, "", 0
+        
     formatted_chat_history = format_chat_history(message, history)
     # Generate response using QA chain
     response = qa_chain.invoke({"question": message, "chat_history": formatted_chat_history})
     response_answer = response["answer"]
     if response_answer.find("Helpful Answer:") != -1:
         response_answer = response_answer.split("Helpful Answer:")[-1]
+    
     response_sources = response["source_documents"]
-    response_source1 = response_sources[0].page_content.strip()
-    response_source2 = response_sources[1].page_content.strip()
-    response_source3 = response_sources[2].page_content.strip()
-    # Langchain sources are zero-based
-    response_source1_page = response_sources[0].metadata["page"] + 1
-    response_source2_page = response_sources[1].metadata["page"] + 1
-    response_source3_page = response_sources[2].metadata["page"] + 1
+    
+    # Handle source documents
+    source_contents = ["", "", ""]
+    source_pages = [0, 0, 0]
+    
+    for i, source in enumerate(response_sources[:3]):
+        source_contents[i] = source.page_content.strip()
+        # Check if the metadata contains a page number
+        if "page" in source.metadata:
+            source_pages[i] = source.metadata["page"] + 1
+        elif "source" in source.metadata:
+            source_pages[i] = 1
+            source_contents[i] = f"From: {source.metadata['source']}\n{source_contents[i]}"
+    
     # Append user message and response to chat history
     new_history = history + [(message, response_answer)]
-    return qa_chain, gr.update(value=""), new_history, response_source1, response_source1_page, response_source2, response_source2_page, response_source3, response_source3_page
     
+    return qa_chain, gr.update(value=""), new_history, source_contents[0], source_pages[0], source_contents[1], source_pages[1], source_contents[2], source_pages[2]
+    
+
+def get_file_icon(file_path):
+    """Return an appropriate emoji icon based on file extension"""
+    ext = Path(file_path).suffix.lower()
+    icons = {
+        '.pdf': '📄',
+        '.txt': '📝',
+        '.md': '📋',
+        '.py': '🐍',
+        '.js': '⚙️',
+        '.json': '📊',
+        '.jsonl': '📊',
+        '.csv': '📈',
+        '.html': '🌐',
+        '.css': '🎨',
+    }
+    return icons.get(ext, '📁')
 
-def upload_file(file_obj):
-    list_file_path = []
-    for idx, file in enumerate(file_obj):
-        file_path = file_obj.name
-        list_file_path.append(file_path)
-    return list_file_path
+
+def display_file_list(file_obj):
+    if not file_obj:
+        return "No files uploaded yet"
+    
+    file_list = [f"{get_file_icon(x.name)} {Path(x.name).name}" for x in file_obj if x is not None]
+    return "\n".join(file_list)
 
 
 def demo():
-    # with gr.Blocks(theme=gr.themes.Default(primary_hue="sky")) as demo:
-    with gr.Blocks(theme=gr.themes.Default(primary_hue="red", secondary_hue="pink", neutral_hue = "sky")) as demo:
+    with gr.Blocks(theme=gr.themes.Soft(primary_hue="indigo", secondary_hue="blue", neutral_hue="sky")) as demo:
         vector_db = gr.State()
         qa_chain = gr.State()
-        gr.HTML("<center><h1>RAG PDF chatbot</h1><center>")
-        gr.Markdown("""<b>Query your PDF documents!</b> This AI agent is designed to perform retrieval augmented generation (RAG) on PDF documents. The app is hosted on Hugging Face Hub for the sole purpose of demonstration. \
-        <b>Please do not upload confidential documents.</b>
+        
+        gr.HTML("<center><h1>📚 Enhanced RAG Chatbot</h1></center>")
+        gr.Markdown("""<b>Query your documents!</b> This enhanced AI agent performs retrieval augmented generation (RAG) on various document types 
+        including PDFs, text files, markdown, code files, and structured data (CSV, JSON, JSONL). <b>Please do not upload confidential documents.</b>
         """)
+        
         with gr.Row():
-            with gr.Column(scale = 86):
-                gr.Markdown("<b>Step 1 - Upload PDF documents and Initialize RAG pipeline</b>")
+            with gr.Column(scale=86):
+                gr.Markdown("<b>Step 1 - Upload Documents and Initialize RAG Pipeline</b>")
                 with gr.Row():
-                    document = gr.Files(height=300, file_count="multiple", file_types=["pdf"], interactive=True, label="Upload PDF documents")
+                    with gr.Column(scale=7):
+                        document = gr.Files(
+                            height=300, 
+                            file_count="multiple", 
+                            file_types=[".pdf", ".txt", ".md", ".py", ".js", ".json", ".jsonl", ".csv", ".html", ".css"], 
+                            interactive=True, 
+                            label="Upload Documents"
+                        )
+                    with gr.Column(scale=3):
+                        file_list = gr.Textbox(
+                            label="Uploaded Files", 
+                            value="No files uploaded yet",
+                            interactive=False,
+                            lines=12
+                        )
+                        document.upload(
+                            display_file_list, 
+                            inputs=[document], 
+                            outputs=[file_list]
+                        )
+                    
                 with gr.Row():
-                    db_btn = gr.Button("Create vector database")
+                    db_btn = gr.Button("Create Vector Database", variant="primary")
+                    
                 with gr.Row():
-                        db_progress = gr.Textbox(value="Not initialized", show_label=False) # label="Vector database status", 
-                gr.Markdown("<style>body { font-size: 16px; }</style><b>Select Large Language Model (LLM) and input parameters</b>")
+                    db_progress = gr.Textbox(
+                        value="Not initialized", 
+                        show_label=False,
+                        container=True
+                    )
+                    
+                gr.Markdown("<b>Step 2 - Select LLM and Parameters</b>")
                 with gr.Row():
-                    llm_btn = gr.Radio(list_llm_simple, label="Available LLMs", value = list_llm_simple[0], type="index") # info="Select LLM", show_label=False
+                    llm_btn = gr.Radio(
+                        list_llm_simple, 
+                        label="Available LLMs", 
+                        value=list_llm_simple[0], 
+                        type="index"
+                    )
+                    
                 with gr.Row():
-                    with gr.Accordion("LLM input parameters", open=False):
+                    with gr.Accordion("LLM Parameters", open=False):
                         with gr.Row():
-                            slider_temperature = gr.Slider(minimum = 0.01, maximum = 1.0, value=0.5, step=0.1, label="Temperature", info="Controls randomness in token generation", interactive=True)
+                            slider_temperature = gr.Slider(
+                                minimum=0.01, 
+                                maximum=1.0, 
+                                value=0.5, 
+                                step=0.1, 
+                                label="Temperature", 
+                                info="Controls randomness in generation", 
+                                interactive=True
+                            )
                         with gr.Row():
-                            slider_maxtokens = gr.Slider(minimum = 128, maximum = 9192, value=4096, step=128, label="Max New Tokens", info="Maximum number of tokens to be generated",interactive=True)
+                            slider_maxtokens = gr.Slider(
+                                minimum=128, 
+                                maximum=9192, 
+                                value=4096, 
+                                step=128, 
+                                label="Max New Tokens", 
+                                info="Maximum tokens to generate",
+                                interactive=True
+                            )
                         with gr.Row():
-                                slider_topk = gr.Slider(minimum = 1, maximum = 10, value=3, step=1, label="top-k", info="Number of tokens to select the next token from", interactive=True)
+                            slider_topk = gr.Slider(
+                                minimum=1, 
+                                maximum=10, 
+                                value=3, 
+                                step=1, 
+                                label="Top-k", 
+                                info="Number of tokens to consider", 
+                                interactive=True
+                            )
+                            
                 with gr.Row():
-                    qachain_btn = gr.Button("Initialize Question Answering Chatbot")
+                    qachain_btn = gr.Button("Initialize Chatbot", variant="primary")
+                    
                 with gr.Row():
-                        llm_progress = gr.Textbox(value="Not initialized", show_label=False) # label="Chatbot status", 
+                    llm_progress = gr.Textbox(
+                        value="Not initialized", 
+                        show_label=False,
+                        container=True
+                    )
 
-            with gr.Column(scale = 200):
-                gr.Markdown("<b>Step 2 - Chat with your Document</b>")
+            with gr.Column(scale=200):
+                gr.Markdown("<b>Step 3 - Chat with Your Documents</b>")
                 chatbot = gr.Chatbot(height=505)
-                with gr.Accordion("Relevent context from the source document", open=False):
+                
+                with gr.Accordion("Relevant Context from Documents", open=False):
                     with gr.Row():
                         doc_source1 = gr.Textbox(label="Reference 1", lines=2, container=True, scale=20)
                         source1_page = gr.Number(label="Page", scale=1)
@@ -184,36 +375,58 @@ def demo():
                     with gr.Row():
                         doc_source3 = gr.Textbox(label="Reference 3", lines=2, container=True, scale=20)
                         source3_page = gr.Number(label="Page", scale=1)
+                        
                 with gr.Row():
-                    msg = gr.Textbox(placeholder="Ask a question", container=True)
+                    msg = gr.Textbox(
+                        placeholder="Ask a question about your documents...", 
+                        container=True,
+                        lines=2
+                    )
+                    
                 with gr.Row():
-                    submit_btn = gr.Button("Submit")
+                    submit_btn = gr.Button("Submit", variant="primary")
                     clear_btn = gr.ClearButton([msg, chatbot], value="Clear")
             
         # Preprocessing events
-        db_btn.click(initialize_database, \
-            inputs=[document], \
-            outputs=[vector_db, db_progress])
-        qachain_btn.click(initialize_LLM, \
-            inputs=[llm_btn, slider_temperature, slider_maxtokens, slider_topk, vector_db], \
-            outputs=[qa_chain, llm_progress]).then(lambda:[None,"",0,"",0,"",0], \
-            inputs=None, \
-            outputs=[chatbot, doc_source1, source1_page, doc_source2, source2_page, doc_source3, source3_page], \
-            queue=False)
+        db_btn.click(
+            initialize_database, 
+            inputs=[document], 
+            outputs=[vector_db, db_progress]
+        )
+        
+        qachain_btn.click(
+            initialize_LLM, 
+            inputs=[llm_btn, slider_temperature, slider_maxtokens, slider_topk, vector_db], 
+            outputs=[qa_chain, llm_progress]
+        ).then(
+            lambda:[None,"",0,"",0,"",0], 
+            inputs=None, 
+            outputs=[chatbot, doc_source1, source1_page, doc_source2, source2_page, doc_source3, source3_page], 
+            queue=False
+        )
 
         # Chatbot events
-        msg.submit(conversation, \
-            inputs=[qa_chain, msg, chatbot], \
-            outputs=[qa_chain, msg, chatbot, doc_source1, source1_page, doc_source2, source2_page, doc_source3, source3_page], \
-            queue=False)
-        submit_btn.click(conversation, \
-            inputs=[qa_chain, msg, chatbot], \
-            outputs=[qa_chain, msg, chatbot, doc_source1, source1_page, doc_source2, source2_page, doc_source3, source3_page], \
-            queue=False)
-        clear_btn.click(lambda:[None,"",0,"",0,"",0], \
-            inputs=None, \
-            outputs=[chatbot, doc_source1, source1_page, doc_source2, source2_page, doc_source3, source3_page], \
-            queue=False)
+        msg.submit(
+            conversation, 
+            inputs=[qa_chain, msg, chatbot], 
+            outputs=[qa_chain, msg, chatbot, doc_source1, source1_page, doc_source2, source2_page, doc_source3, source3_page], 
+            queue=False
+        )
+        
+        submit_btn.click(
+            conversation, 
+            inputs=[qa_chain, msg, chatbot], 
+            outputs=[qa_chain, msg, chatbot, doc_source1, source1_page, doc_source2, source2_page, doc_source3, source3_page], 
+            queue=False
+        )
+        
+        clear_btn.click(
+            lambda:[None,"",0,"",0,"",0], 
+            inputs=None, 
+            outputs=[chatbot, doc_source1, source1_page, doc_source2, source2_page, doc_source3, source3_page], 
+            queue=False
+        )
+        
     demo.queue().launch(debug=True)