Create app.py

app.py

@@ -0,0 +1,211 @@
+import gradio as gr
+import os
+import time
+from datetime import datetime
+from langchain_community.document_loaders import PyPDFLoader, TextLoader, Docx2txtLoader
+from langchain_community.vectorstores import Chroma
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_huggingface import HuggingFaceEndpoint, HuggingFaceEmbeddings
+from langchain.chains import ConversationalRetrievalChain
+from langchain.memory import ConversationBufferMemory
+from pptx import Presentation
+from io import BytesIO
+
+# Environment setup for Hugging Face token
+os.environ["HUGGINGFACEHUB_API_TOKEN"] = os.getenv("HUGGINGFACEHUB_API_TOKEN", "your-hf-token-here")
+
+# Model and embedding options
+LLM_MODELS = {
+    "Lightweight (Gemma-2B)": "google/gemma-2b-it",
+    "Balanced (Mixtral-8x7B)": "mistralai/Mixtral-8x7B-Instruct-v0.1",
+    "High Accuracy (Llama-3-8B)": "meta-llama/Llama-3-8b-hf"
+}
+
+EMBEDDING_MODELS = {
+    "Lightweight (MiniLM-L6)": "sentence-transformers/all-MiniLM-L6-v2",
+    "Balanced (MPNet-Base)": "sentence-transformers/all-mpnet-base-v2",
+    "High Accuracy (BGE-Large)": "BAAI/bge-large-en-v1.5"
+}
+
+# Global state
+vector_store = None
+qa_chain = None
+chat_history = []
+memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
+
+# Custom PPTX loader
+class PPTXLoader:
+    def __init__(self, file_path):
+        self.file_path = file_path
+
+    def load(self):
+        docs = []
+        with open(self.file_path, "rb") as f:
+            prs = Presentation(BytesIO(f.read()))
+            for slide_num, slide in enumerate(prs.slides, 1):
+                text = ""
+                for shape in slide.shapes:
+                    if hasattr(shape, "text"):
+                        text += shape.text + "\n"
+                if text.strip():
+                    docs.append({"page_content": text, "metadata": {"source": self.file_path, "slide": slide_num}})
+        return docs
+
+# Function to load documents
+def load_documents(files):
+    documents = []
+    for file in files:
+        file_path = file.name
+        if file_path.endswith(".pdf"):
+            loader = PyPDFLoader(file_path)
+            documents.extend(loader.load())
+        elif file_path.endswith(".txt"):
+            loader = TextLoader(file_path)
+            documents.extend(loader.load())
+        elif file_path.endswith(".docx"):
+            loader = Docx2txtLoader(file_path)
+            documents.extend(loader.load())
+        elif file_path.endswith(".pptx"):
+            loader = PPTXLoader(file_path)
+            documents.extend([{"page_content": doc["page_content"], "metadata": doc["metadata"]} for doc in loader.load()])
+    return documents
+
+# Function to process documents and create vector store
+def process_documents(files, chunk_size, chunk_overlap, embedding_model):
+    global vector_store, qa_chain
+    if not files:
+        return "Please upload at least one document.", None
+
+    # Load documents
+    documents = load_documents(files)
+    if not documents:
+        return "No valid documents loaded.", None
+
+    # Split documents
+    text_splitter = RecursiveCharacterTextSplitter(
+        chunk_size=int(chunk_size),
+        chunk_overlap=int(chunk_overlap),
+        length_function=len
+    )
+    doc_splits = text_splitter.split_documents(documents)
+
+    # Create embeddings
+    embeddings = HuggingFaceEmbeddings(model_name=EMBEDDING_MODELS[embedding_model])
+
+    # Create vector store
+    try:
+        vector_store = Chroma.from_documents(doc_splits, embeddings, persist_directory="./chroma_db")
+        return f"Processed {len(documents)} documents into {len(doc_splits)} chunks.", None
+    except Exception as e:
+        return f"Error processing documents: {str(e)}", None
+
+# Function to initialize QA chain
+def initialize_qa_chain(llm_model, temperature):
+    global qa_chain
+    try:
+        llm = HuggingFaceEndpoint(
+            repo_id=LLM_MODELS[llm_model],
+            temperature=float(temperature),
+            max_length=512,
+            huggingfacehub_api_token=os.environ["HUGGINGFACEHUB_API_TOKEN"]
+        )
+        qa_chain = ConversationalRetrievalChain.from_llm(
+            llm=llm,
+            retriever=vector_store.as_retriever(search_kwargs={"k": 3}),
+            memory=memory
+        )
+        return "QA chain initialized successfully.", None
+    except Exception as e:
+        return f"Error initializing QA chain: {str(e)}", None
+
+# Function to handle user query
+def answer_question(question, llm_model, embedding_model, temperature, chunk_size, chunk_overlap):
+    global chat_history
+    if not vector_store or not qa_chain:
+        return "Please upload documents and initialize the QA chain.", chat_history
+
+    try:
+        response = qa_chain({"question": question})["answer"]
+        chat_history.append(("User", question))
+        chat_history.append(("Bot", response))
+        return response, chat_history
+    except Exception as e:
+        return f"Error answering question: {str(e)}", chat_history
+
+# Function to export chat history
+def export_chat():
+    if not chat_history:
+        return "No chat history to export.", None
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    filename = f"chat_history_{timestamp}.txt"
+    with open(filename, "w") as f:
+        for role, message in chat_history:
+            f.write(f"{role}: {message}\n\n")
+    return f"Chat history exported to {filename}.", filename
+
+# Function to reset the app
+def reset_app():
+    global vector_store, qa_chain, chat_history, memory
+    vector_store = None
+    qa_chain = None
+    chat_history = []
+    memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
+    if os.path.exists("./chroma_db"):
+        import shutil
+        shutil.rmtree("./chroma_db")
+    return "App reset successfully.", None
+
+# Gradio interface
+with gr.Blocks(theme=gr.themes.Soft(), title="DocTalk: Document Q&A Chatbot") as demo:
+    gr.Markdown("# DocTalk: Document Q&A Chatbot")
+    gr.Markdown("Upload documents (PDF, TXT, DOCX, PPTX), select models, tune parameters, and ask questions!")
+
+    with gr.Row():
+        with gr.Column(scale=2):
+            file_upload = gr.Files(label="Upload Documents", file_types=[".pdf", ".txt", ".docx", ".pptx"])
+            with gr.Row():
+                process_button = gr.Button("Process Documents")
+                reset_button = gr.Button("Reset App")
+            status = gr.Textbox(label="Status", interactive=False)
+
+        with gr.Column(scale=1):
+            llm_model = gr.Dropdown(choices=list(LLM_MODELS.keys()), label="Select LLM Model", value="Lightweight (Gemma-2B)")
+            embedding_model = gr.Dropdown(choices=list(EMBEDDING_MODELS.keys()), label="Select Embedding Model", value="Lightweight (MiniLM-L6)")
+            temperature = gr.Slider(minimum=0.0, maximum=1.0, step=0.1, value=0.7, label="Temperature")
+            chunk_size = gr.Slider(minimum=500, maximum=2000, step=100, value=1000, label="Chunk Size")
+            chunk_overlap = gr.Slider(minimum=0, maximum=500, step=50, value=100, label="Chunk Overlap")
+            init_button = gr.Button("Initialize QA Chain")
+
+    gr.Markdown("## Chat Interface")
+    question = gr.Textbox(label="Ask a Question", placeholder="Type your question here...")
+    answer = gr.Textbox(label="Answer", interactive=False)
+    chat_display = gr.Chatbot(label="Chat History")
+    export_button = gr.Button("Export Chat History")
+    export_file = gr.File(label="Exported Chat File")
+
+    # Event handlers
+    process_button.click(
+        fn=process_documents,
+        inputs=[file_upload, chunk_size, chunk_overlap, embedding_model],
+        outputs=[status, chat_display]
+    )
+    init_button.click(
+        fn=initialize_qa_chain,
+        inputs=[llm_model, temperature],
+        outputs=[status, chat_display]
+    )
+    question.submit(
+        fn=answer_question,
+        inputs=[question, llm_model, embedding_model, temperature, chunk_size, chunk_overlap],
+        outputs=[answer, chat_display]
+    )
+    export_button.click(
+        fn=export_chat,
+        outputs=[status, export_file]
+    )
+    reset_button.click(
+        fn=reset_app,
+        outputs=[status, chat_display]
+    )
+
+demo.launch()