Update app.py

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 import os
-import torch
+from concurrent.futures import ThreadPoolExecutor
 from langchain_community.vectorstores import Chroma
 from langchain_community.document_loaders import PyPDFLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
@@ -12,42 +12,50 @@ from langchain_community.retrievers import BM25Retriever
 from langchain.retrievers import EnsembleRetriever
 
 # Environment variable for API token
-api_token = os.getenv("API_TOKEN")
-print(f"API Token loaded: {api_token[:5]}...")  # Debug: Show first 5 chars of token
+api_token = os.getenv("FirstToken")
+print(f"API Token loaded: {api_token[:5]}...")  # Debug
 if not api_token:
-    raise ValueError("Environment variable 'FirstToken' not set. Please set the Hugging Face API token.")
+    raise ValueError("Environment variable 'FirstToken' not set.")
 
 # Available LLM models
 list_llm = [
-    "mistralai/Mixtral-8x7B-Instruct-v0.1",  # Publicly accessible
+    "mistralai/Mixtral-8x7B-Instruct-v0.1",
     "mistralai/Mistral-7B-Instruct-v0.2",
     "deepseek-ai/deepseek-llm-7b-chat"
 ]
 list_llm_simple = [os.path.basename(llm) for llm in list_llm]
 
 # -----------------------------------------------------------------------------
-# Document Loading and Splitting
+# Document Loading and Splitting (Optimized with Threading)
 # -----------------------------------------------------------------------------
+def load_single_pdf(file_path):
+    """Load a single PDF file."""
+    loader = PyPDFLoader(file_path)
+    return loader.load()
+
 def load_doc(list_file_path, progress=gr.Progress()):
-    """Load and split PDF documents into chunks."""
+    """Load and split PDF documents into chunks with multi-threading."""
     if not list_file_path:
         raise ValueError("No files provided for processing.")
     
-    loaders = [PyPDFLoader(x) for x in list_file_path]
-    pages = []
-    for i, loader in enumerate(loaders):
-        progress((i + 1) / len(loaders), "Loading PDFs...")
-        pages.extend(loader.load())
+    # Use ThreadPoolExecutor to parallelize PDF loading
+    with ThreadPoolExecutor() as executor:
+        pages = list(executor.map(load_single_pdf, list_file_path))
+        pages = [page for sublist in pages for page in sublist]  # Flatten list
     
-    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1024, chunk_overlap=64)
-    return text_splitter.split_documents(pages)
+    progress(0.5, "Splitting documents...")
+    text_splitter = RecursiveCharacterTextSplitter(chunk_size=2048, chunk_overlap=128)  # Increased chunk size
+    doc_splits = text_splitter.split_documents(pages)
+    return doc_splits
 
 # -----------------------------------------------------------------------------
-# Vector Database Creation
+# Vector Database Creation (Optimized with Lightweight Embeddings)
 # -----------------------------------------------------------------------------
-def create_chromadb(splits, persist_directory="chroma_db"):
-    """Create ChromaDB vector database from document splits."""
-    embeddings = HuggingFaceEmbeddings()
+def create_chromadb(splits, persist_directory="chroma_db", progress=gr.Progress()):
+    """Create ChromaDB vector database with optimized embeddings."""
+    # Use a lighter embedding model
+    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+    progress(0.7, "Creating vector database...")
     chromadb = Chroma.from_documents(
         documents=splits,
         embedding=embeddings,
@@ -61,13 +69,13 @@ def create_chromadb(splits, persist_directory="chroma_db"):
 def create_bm25_retriever(splits):
     """Create BM25 retriever from document splits."""
     retriever = BM25Retriever.from_documents(splits)
-    retriever.k = 3
+    retriever.k = 2  # Reduced to 2 documents for faster retrieval
     return retriever
 
 def create_ensemble_retriever(vector_db, bm25_retriever):
-    """Create an ensemble retriever combining vector DB and BM25."""
+    """Create an ensemble retriever."""
     return EnsembleRetriever(
-        retrievers=[vector_db.as_retriever(), bm25_retriever],
+        retrievers=[vector_db.as_retriever(search_kwargs={"k": 2}), bm25_retriever],  # Limit to 2 docs
         weights=[0.7, 0.3]
     )
 
@@ -78,10 +86,12 @@ def initialize_database(list_file_obj, progress=gr.Progress()):
     """Initialize the document database with error handling."""
     try:
         list_file_path = [x.name for x in list_file_obj if x is not None]
+        progress(0.1, "Loading documents...")
         doc_splits = load_doc(list_file_path, progress)
-        chromadb = create_chromadb(doc_splits)
+        chromadb = create_chromadb(doc_splits, progress=progress)
         bm25_retriever = create_bm25_retriever(doc_splits)
         ensemble_retriever = create_ensemble_retriever(chromadb, bm25_retriever)
+        progress(1.0, "Database creation complete!")
         return ensemble_retriever, "Database created successfully!"
     except Exception as e:
         return None, f"Error initializing database: {str(e)}"
@@ -90,7 +100,7 @@ def initialize_database(list_file_obj, progress=gr.Progress()):
 # Initialize LLM Chain
 # -----------------------------------------------------------------------------
 def initialize_llmchain(llm_model, temperature, max_tokens, top_k, retriever):
-    """Initialize the language model chain with error handling."""
+    """Initialize the language model chain."""
     if retriever is None:
         raise ValueError("Retriever is None. Please process documents first.")
     
@@ -211,9 +221,9 @@ def demo():
                     language_btn = gr.Radio(choices=["English", "Português"], label="Response Language", value="English")
                     with gr.Accordion("Advanced Settings", open=False):
                         slider_temperature = gr.Slider(0.01, 1.0, value=0.5, step=0.1, label="Analysis Precision")
-                        slider_maxtokens = gr.Slider(128, 9192, value=4096, step=128, label="Response Length")
-                        slider_topk = gr.Slider(1, 10, value=3, step=1, label="Analysis Diversity")
-                    qachain_btn = gr.Button("Initialize Assistant", interactive=False)  # Disabled by default
+                        slider_maxtokens = gr.Slider(128, 2048, value=1024, step=128, label="Response Length")  # Reduced max_tokens
+                        slider_topk = gr.Slider(1, 5, value=3, step=1, label="Analysis Diversity")  # Reduced range
+                    qachain_btn = gr.Button("Initialize Assistant", interactive=False)
                     llm_progress = gr.Textbox(value="Not initialized", label="Assistant Status")
 
             with gr.Column(scale=2):