Uploaded RAG UI gradio app

app.py

@@ -1,64 +1,562 @@
 import gradio as gr
-from huggingface_hub import InferenceClient
-
-"""
-For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
-"""
-client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
-
-
-def respond(
-    message,
-    history: list[tuple[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-):
-    messages = [{"role": "system", "content": system_message}]
-
-    for val in history:
-        if val[0]:
-            messages.append({"role": "user", "content": val[0]})
-        if val[1]:
-            messages.append({"role": "assistant", "content": val[1]})
-
-    messages.append({"role": "user", "content": message})
-
-    response = ""
-
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-    ):
-        token = message.choices[0].delta.content
-
-        response += token
-        yield response
-
-
-"""
-For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
-"""
-demo = gr.ChatInterface(
-    respond,
-    additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",
-        ),
-    ],
+import csv
+import random
+import os
+import shutil
+import json
+from llama_index.embeddings.openai import OpenAIEmbedding
+from llama_index.core import (
+    VectorStoreIndex,
+    SimpleDirectoryReader,
+    StorageContext,
+    load_index_from_storage,
 )
+from llama_index.core.settings import Settings
+import faiss
+import numpy as np
+from llama_index.vector_stores.faiss import FaissVectorStore
+from llama_index.core.node_parser import SimpleNodeParser, SentenceSplitter
+from llama_index.core.schema import Document
+from llama_index.core.schema import IndexNode
+from llama_index.core import ServiceContext
+from llama_index.core.query_engine.retriever_query_engine import RetrieverQueryEngine
+from llama_index.embeddings.huggingface.base import HuggingFaceEmbedding
+from llama_index.llms.openai import OpenAI
+from transformers import BitsAndBytesConfig
+from llama_index.core.prompts import PromptTemplate
+import torch
+import pandas as pd
+import fitz
+from transformers import pipeline
+from sklearn.metrics.pairwise import cosine_similarity
+from sklearn.feature_extraction.text import TfidfVectorizer
 
 
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+# os.environ["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY")
+llm = OpenAI(temperature=0, model="gpt-4o-mini", max_tokens=512)
+Settings.llm = llm
+
+UPLOAD_DIR = "uploaded_files"
+STATE_FILE = "uploaded_files_state.json"
+PERSIST_DIR = "persisted_indexes"
+
+os.makedirs(UPLOAD_DIR, exist_ok=True)
+os.makedirs(PERSIST_DIR, exist_ok=True)
+
+# !!! why???
+# torch.set_num_threads(1)
+# torch.set_num_interop_threads(1)
+
+
+def index_gen(file_path, index_name):
+    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+
+    # One giant index: insertion example
+    # if os.path.exists('persisted_indexes/test1.faiss'):
+    #     print("RUNNING TEST!")
+    #     # Load document from file
+    #     documents = SimpleDirectoryReader(input_files=[file_path]).load_data()
+
+    #     faiss_index = faiss.read_index('persisted_indexes/test1.faiss')
+    #     embed_model = HuggingFaceEmbedding(
+    #         model_name="BAAI/bge-small-en-v1.5"
+    #     )
+    #     Settings.embed_model = embed_model
+
+    #     vector_store = FaissVectorStore(faiss_index=faiss_index)
+    #     storage_context = StorageContext.from_defaults(
+    #         persist_dir=PERSIST_DIR, vector_store=vector_store
+    #     )
+
+    #     index = load_index_from_storage(storage_context)
+    #     print(index)
+    #     for doc in documents:
+    #         print('inserting ', doc)
+    #         index.insert(doc)
+    #     index.storage_context.persist(PERSIST_DIR)
+    #     faiss.write_index(faiss_index, 'persisted_indexes/test1.faiss')
+    #     print('insertion and persist complete!')
+    #     return index
+
+
+    try:
+        # Load document from file
+        documents = SimpleDirectoryReader(input_files=[file_path]).load_data()
+
+        # Initialize embedding model and vector store
+        embed_model = HuggingFaceEmbedding(
+            model_name="BAAI/bge-small-en-v1.5", device=device
+        )
+        Settings.embed_model = embed_model
+        embedding_dim = 384  # Ensure this matches the embedding model used
+
+        faiss_index = faiss.IndexFlatL2(embedding_dim)
+        vector_store = FaissVectorStore(faiss_index=faiss_index)
+        storage_context = StorageContext.from_defaults(vector_store=vector_store)
+
+        print(f"Number of documents to index: {len(documents)}.")
+
+        # Parse and index documents
+        parser = SentenceSplitter()
+        nodes = parser.get_nodes_from_documents(documents)
+        index = VectorStoreIndex(nodes, storage_context=storage_context)
+        print(f"Number of nodes generated:{len(nodes)}")
+
+        # individual index directory
+        index_directory = os.path.join(PERSIST_DIR, index_name)
+        os.makedirs(index_directory, exist_ok=True)
+        index_path = os.path.join(index_directory, f"{index_name}.faiss")
+
+
+        index.storage_context.persist(index_directory)
+        # index.storage_context.persist(PERSIST_DIR)
+        faiss.write_index(faiss_index, index_path)
+
+        if not os.path.exists(index_path):
+            raise FileNotFoundError(
+                f"FAISS index file not created at path: {index_path}"
+            )
+
+        return index_path
+
+    except Exception as e:
+        print(f"Error in index_gen with file {file_path}: {str(e)}")
+        return None
+
+
+def save_uploaded_files_state(uploaded_files, indexed_files=None):
+    try:
+        state_file_json = {}
+        state_file_json["uploaded_files"] = list(uploaded_files)
+
+        if indexed_files:
+            state_file_json["indexed_files"] = list(indexed_files)
+
+        # else:
+        #     # ??? why
+        #     _, existing_indexed_files = load_uploaded_files_state()
+        #     state_file_json["indexed_files"] = list(existing_indexed_files)
+
+        with open(STATE_FILE, "w") as f:
+            json.dump(state_file_json, f, indent=4)
+
+    except IOError as e:
+        print(f"Error saving uploaded files state: {str(e)}")
+
+
+def load_uploaded_files_state():
+    try:
+        if os.path.exists(STATE_FILE):
+            with open(STATE_FILE, "r") as f:
+                state_data = json.load(f)
+                return set(state_data.get("uploaded_files", set())), set(
+                    state_data.get("indexed_files", set())
+                )
+
+    except (IOError, json.JSONDecodeError) as e:
+        print(f"Error loading uploaded files state: {str(e)}")
+
+    return set(), set()
+
+
+def save_file(file_path):
+    try:
+        file_name = os.path.basename(file_path)
+        server_save_path = os.path.join(UPLOAD_DIR, file_name)
+        shutil.copy(file_path, server_save_path)
+        return server_save_path
+
+    except (IOError, shutil.Error) as e:
+        print(f"Error saving file {file_path}: {str(e)}")
+        return None
+
+
+with gr.Blocks() as demo:
+    gr.Markdown("## 📁 File Management & Chat Assistant")
+
+    with gr.Tabs():
+        # Tab 1: File Management
+        with gr.Tab("File Management"):
+            with gr.Row():
+                with gr.Column(scale=1):
+                    file_upload = gr.File(
+                        label="Upload PDF,JSON or TXT Files",
+                        file_types=[".pdf", ".json", ".txt", "directory"],
+                        file_count="multiple",
+                        interactive=True,
+                    )
+                    file_table = gr.DataFrame(
+                        headers=["Sr. No.", "File Name", "File Size"],
+                        value=[],
+                        interactive=False,
+                        row_count=(4, "dynamic"),
+                        wrap=True,
+                        max_height=1000
+                    )
+                    file_checkbox = gr.CheckboxGroup(
+                        label="Select Files to Index/Delete", choices=[]
+                    )
+                    select_all_button = gr.Button("Select All")
+                    index_button = gr.Button("Index Selected Files")
+                    delete_button = gr.Button("Delete Selected Files")
+
+                with gr.Column(scale=3):
+                    message_box = gr.Markdown("")
+                    chatbot = gr.Chatbot(label="LLM", type="messages")
+
+                    with gr.Row():
+                        chat_input = gr.Textbox(
+                            show_label=False,
+                            placeholder="Type your message here",
+                            scale=8,
+                        )
+                        send_button = gr.Button("Send", scale=1)
+
+        # Tab 2: Indexed Files
+        with gr.Tab("Indexed Files"):
+            indexed_file_table = gr.DataFrame(
+                headers=["Indexed File", "Size"],
+                value=[],
+                interactive=False,
+                row_count=(4, "dynamic"),
+            )
+
+    # STATES
+    uploaded_files_state = gr.State(load_uploaded_files_state())
+
+    @delete_button.click(
+        inputs=[file_checkbox, uploaded_files_state, file_upload],
+        outputs=[file_table, file_checkbox, uploaded_files_state, indexed_file_table],
+    )
+    def delete_files(selected_files, uploaded_files_state, file_upload):
+        print("deleting files...: ", selected_files, uploaded_files_state, file_upload)
+
+        uploaded_files, indexed_files = uploaded_files_state
+
+        if not selected_files or not uploaded_files:
+            return gr.update(), selected_files, (uploaded_files, indexed_files)
+
+        # default return
+        # return [[]], selected_files, uploaded_files_state
+
+        # "we" means with extension
+        selected_file_names_we = [file.split(". ")[1] for file in selected_files]
+
+        for file_name_we in selected_file_names_we:
+            file_path = os.path.join(UPLOAD_DIR, file_name_we)
+            index_name = file_name_we.split(".")[0]
+            index_directory = os.path.join(PERSIST_DIR, index_name)
+            index_path = os.path.join(index_directory, f'{index_name}.faiss')
+            print(file_name_we, file_path, index_name, index_directory, index_path)
+
+            try:
+                if os.path.exists(file_path):
+                    os.remove(file_path)
+                    uploaded_files.remove(file_path)
+
+                else:
+                    gr.Error(f"Could not delete file (File not found): {file_path}", duration=3)
+
+                if os.path.exists(index_directory):
+                    shutil.rmtree(index_directory)
+                    indexed_files.remove(index_path)
+
+                else:
+                    gr.Error(f"Could not delete index directory (Path not found): {index_directory}", duration=3)
+
+            except Exception as e:
+                gr.Error(f"Error deleting {file_name_we}: {str(e)}", duration=3)
+
+        save_uploaded_files_state(uploaded_files, indexed_files)
+
+        file_info, checkbox_options = [], []
+        for idx, file_path in enumerate(uploaded_files, start=1):
+            file_name = os.path.basename(file_path)
+            file_size = os.path.getsize(file_path)
+            file_info.append([idx, file_name, f"{round(file_size / 1024, 2)} KB"])
+            checkbox_options.append(f"{idx}. {file_name}")
+
+        indexed_file_display = [
+            [
+                os.path.basename(index_path).split(".")[0],
+                f"{round(os.path.getsize(index_path) / 1024, 2)} KB",
+            ]
+            for index_path in indexed_files
+        ]
+
+        return (
+            file_info,
+            gr.update(choices=checkbox_options, value=[]),
+            (uploaded_files, indexed_files),
+            indexed_file_display,
+        )
+
+    @chat_input.submit(
+        inputs=[chat_input, chatbot, uploaded_files_state],
+        outputs=[chat_input, chatbot],
+    )
+    @send_button.click(
+        inputs=[chat_input, chatbot, uploaded_files_state],
+        outputs=[chat_input, chatbot],
+    )
+    # Chat function with improved SQuAD matching
+    def chat_with_bot(user_input, chat_history, uploaded_files_state):
+        if not user_input:
+            return user_input, chat_history
+
+        _, indexed_files = uploaded_files_state
+
+        chat_history.append(
+            {
+                "role": "user",
+                "content": user_input,
+            }
+        )
+
+        response = "I do not have the answer. Please upload and index relevant files first."
+        file_with_answer = None
+        custom_prompt = PromptTemplate(
+            template=(
+                "Use the following context to answer the query. Do not use outside knowledge. "
+                "If the answer is not found in the context, respond with: 'I do not have the answer.'\n"
+                "Context: {context_str}\n"
+                "Query: {query_str}\n"
+                "Answer:"
+            )
+        )
+
+        if not index_files:
+            response = "No files have been indexed for answering this question."
+
+        try:
+            for index_path in indexed_files:
+                print('checking ', index_path)
+                file_name = os.path.basename(index_path)
+                index_name = file_name.split(".")[0]
+
+                if not os.path.exists(index_path):
+                    print(f"FAISS index not found at {index_path}, skipping...")
+                    continue
+
+                storage_context = None
+                try:
+                    faiss_index = faiss.read_index(index_path)
+                    embed_model = HuggingFaceEmbedding(
+                        model_name="BAAI/bge-small-en-v1.5"
+                    )
+                    Settings.embed_model = embed_model
+
+                    vector_store = FaissVectorStore(faiss_index=faiss_index)
+                    storage_context = StorageContext.from_defaults(
+                        persist_dir=f'{PERSIST_DIR}/{index_name}', vector_store=vector_store
+                    )
+
+                except Exception as e:
+                    raise RuntimeError(
+                        f"Failed to load FAISS index at {index_path}: {str(e)}"
+                    )
+
+                # print(get_global("embed_model"))
+
+                index = load_index_from_storage(storage_context)
+                print(f"Index loaded with {len(index.docstore.docs)} documents.")
+
+                retriever = index.as_retriever(similarity_top_k=10)
+                query_engine = RetrieverQueryEngine(retriever=retriever)
+                query_engine.update_prompts(
+                    {"response_synthesizer:text_qa_template": custom_prompt}
+                )
+
+                # Query the index for the user input
+                query_result = query_engine.query(user_input)
+                print("query result: ", query_result)
+
+                if query_result.response.strip() != "I do not have the answer.":
+                    response = f"{query_result.response} \n\n Source: {file_name}"
+                    # response = f"Answer from indexed file '{file_name}': {query_result.response}"
+                    file_with_answer = file_name
+                    break
+
+                else:
+                    response = "I do not have the answer."
+
+        except Exception as e:
+            response = f"Error querying the index: {str(e)}"
+            print(response)
+
+        chat_history.append(
+            {
+                "role": "assistant",
+                "content": response,
+            }
+        )
+
+        return gr.update(value=""), chat_history
+
+    @index_button.click(
+        inputs=[file_checkbox, uploaded_files_state, indexed_file_table],
+        outputs=[
+            file_checkbox,
+            uploaded_files_state,
+            indexed_file_table,
+            select_all_button,
+        ],
+    )
+    def index_files(selected_files, uploaded_files_state, indexed_file_table):
+        uploaded_files, indexed_files = uploaded_files_state
+        print("indexing files...", selected_files, uploaded_files_state)
+
+        if not selected_files or not uploaded_files:
+            gr.Warning("Please select or upload files for indexing.", duration=3)
+            return (
+                selected_files,
+                uploaded_files_state,
+                indexed_file_table,
+                gr.update(),
+            )
+
+
+        files_to_index = []
+        for file in selected_files:
+            file_name_we = file.split(". ")[1]
+            file_path = os.path.join(UPLOAD_DIR, file_name_we)
+            index_name = file_name_we.split(".")[0]
+            index_directory = os.path.join(PERSIST_DIR, index_name)
+            index_path = os.path.join(index_directory, f'{index_name}.faiss')
+
+            if index_path not in indexed_files:
+                files_to_index.append(file_path)
+            else:
+                gr.Info(
+                    f"File '{os.path.basename(file_path)}' is already indexed.",
+                    duration=3,
+                )
+
+        for file_path in files_to_index:
+            try:
+                file_name = os.path.basename(file_path)
+                index_name = file_name.split(".")[0]
+                index_path = index_gen(file_path, index_name)
+                gr.Info(f"Successfully indexed: {file_name}", duration=3)
+
+                # Save indexed file info for persistence
+                # index_path = os.path.join(PERSIST_DIR, f"{index_name}.faiss")
+                indexed_files.add(index_path)
+
+            except Exception as e:
+                gr.Error(f"Error indexing {file_path}: {str(e)}", duration=3)
+
+        # Update the state with new indexed files
+        save_uploaded_files_state(uploaded_files, indexed_files)
+
+        # Convert indexed file info to display format
+        indexed_file_display = [
+            [
+                os.path.basename(index_path).split(".")[0],
+                f"{round(os.path.getsize(index_path) / 1024, 2)} KB",
+            ]
+            for index_path in indexed_files
+        ]
+
+        return (
+            gr.update(value=[]),
+            (uploaded_files, indexed_files),
+            indexed_file_display,
+            gr.update(value="Select All"),
+        )
+
+    @select_all_button.click(
+        inputs=[uploaded_files_state, select_all_button, file_checkbox],
+        outputs=[file_checkbox, select_all_button],
+    )
+    def select_all_checkbox(uploaded_files_state, select_all_button, file_checkbox):
+        uploaded_files, _ = uploaded_files_state
+
+        if not uploaded_files:
+            return file_checkbox, select_all_button
+
+        button_value = ""
+        if select_all_button == "Select All":
+            button_value = "Unselect All"
+        else:
+            button_value = "Select All"
+
+        checkbox_options = []
+        if not file_checkbox:
+            checkbox_options = [
+                f"{idx + 1}. {os.path.basename(file)}"
+                for idx, file in enumerate(uploaded_files)
+            ]
+
+        return gr.update(value=checkbox_options), gr.update(value=button_value)
+
+    # Load initial state when app starts
+    @demo.load(
+        inputs=[uploaded_files_state],
+        outputs=[file_table, file_checkbox, uploaded_files_state, indexed_file_table],
+    )
+    def load_state_on_start(uploaded_files_state):
+        uploaded_files, indexed_files = load_uploaded_files_state()
+
+        print("demo loading...", uploaded_files, indexed_files)
+
+        # Populate uploaded files table and checkbox options
+        file_info = []
+        checkbox_options = []
+        for idx, server_file_path in enumerate(uploaded_files, start=1):
+            file_name = os.path.basename(server_file_path)
+            file_size = os.path.getsize(server_file_path)
+            file_info.append([idx, file_name, f"{round(file_size / 1024, 2)} KB"])
+            checkbox_options.append(f"{idx}. {file_name}")
+
+        # Populate indexed files table
+        indexed_file_display = [
+            [
+                os.path.basename(index_path).split(".")[0],
+                f"{round(os.path.getsize(index_path) / 1024, 2)} KB",
+            ]
+            for index_path in indexed_files
+        ]
+
+        return (
+            file_info,
+            gr.update(choices=checkbox_options),
+            (uploaded_files, indexed_files),
+            indexed_file_display,
+        )
+
+    @file_upload.upload(
+        inputs=[file_upload, uploaded_files_state],
+        outputs=[file_table, file_checkbox, file_upload, uploaded_files_state],
+    )
+    def upload_files(file_upload, uploaded_files_state):
+        uploaded_files, indexed_files = uploaded_files_state
+
+        for file_path in file_upload:
+            server_save_path = save_file(file_path)
+            if server_save_path:
+                uploaded_files.add(server_save_path)
+
+        save_uploaded_files_state(uploaded_files)
+
+        file_info = []
+        checkbox_options = []
+        for i, file_path in enumerate(uploaded_files, start=1):
+            file_name = os.path.basename(file_path)
+            file_size = os.path.getsize(file_path)
+            file_info.append([i, file_name, f"{round(file_size / 1024, 2)} KB"])
+            checkbox_options.append(f"{i}. {file_name}")
+
+        gr.Info("Successfully uploaded file(s).", duration=3)
+
+        return (
+            file_info,
+            gr.update(choices=checkbox_options),
+            [],
+            (uploaded_files, indexed_files),
+        )
+
 if __name__ == "__main__":
     demo.launch()