remove unused imports and function, rename functions and fix llmchain init progress

app.py

@@ -2,20 +2,15 @@ import os
 import re
 from pathlib import Path
 
-import accelerate
 import chromadb
 import gradio as gr
-import torch
-import tqdm
-import transformers
-from langchain.chains import ConversationalRetrievalChain, ConversationChain
+from langchain.chains import ConversationalRetrievalChain
 from langchain.memory import ConversationBufferMemory
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain_community.document_loaders import PyPDFLoader
 from langchain_community.embeddings import HuggingFaceEmbeddings
-from langchain_community.llms import HuggingFaceEndpoint, HuggingFacePipeline
+from langchain_community.llms import HuggingFaceEndpoint
 from langchain_community.vectorstores import Chroma
-from transformers import AutoTokenizer
 from unidecode import unidecode
 
 list_llm = [
@@ -31,8 +26,7 @@ list_llm = [
 list_llm_simple = [os.path.basename(llm) for llm in list_llm]
 
 
-# Load PDF document and create doc splits
-def load_doc(list_file_path, chunk_size, chunk_overlap):
+def load_doc_and_create_splits(list_file_path, chunk_size, chunk_overlap):
     # Processing for one document only
     # loader = PyPDFLoader(file_path)
     # pages = loader.load()
@@ -48,8 +42,7 @@ def load_doc(list_file_path, chunk_size, chunk_overlap):
     return doc_splits
 
 
-# Create vector database
-def create_db(splits, collection_name):
+def create_vector_db(splits, collection_name):
     embedding = HuggingFaceEmbeddings()
     new_client = chromadb.EphemeralClient()
     vectordb = Chroma.from_documents(
@@ -61,21 +54,10 @@ def create_db(splits, collection_name):
     return vectordb
 
 
-# Load vector database
-def load_db():
-    embedding = HuggingFaceEmbeddings()
-    vectordb = Chroma(embedding_function=embedding)
-    return vectordb
-
-
-# Initialize langchain LLM chain
 def initialize_llmchain(
     llm_model, temperature, max_tokens, top_k, vector_db, progress=gr.Progress()
 ):
-    progress(0.1, desc="Initializing HF tokenizer...")
-
-    # HuggingFaceHub uses HF inference endpoints
-    progress(0.5, desc="Initializing HF Hub...")
+    progress(0.1, desc="Initializing HF Hub...")
     if llm_model == "mistralai/Mixtral-8x7B-Instruct-v0.1":
         llm = HuggingFaceEndpoint(
             repo_id=llm_model,
@@ -92,14 +74,14 @@ def initialize_llmchain(
             top_k=top_k,
         )
 
-    progress(0.75, desc="Defining buffer memory...")
+    progress(0.6, desc="Defining buffer memory...")
     memory = ConversationBufferMemory(
         memory_key="chat_history", output_key="answer", return_messages=True
     )
     # retriever=vector_db.as_retriever(search_type="similarity", search_kwargs={'k': 3})
     retriever = vector_db.as_retriever()
 
-    progress(0.8, desc="Defining retrieval chain...")
+    progress(0.75, desc="Defining retrieval chain...")
     qa_chain = ConversationalRetrievalChain.from_llm(
         llm,
         retriever=retriever,
@@ -108,6 +90,7 @@ def initialize_llmchain(
         return_source_documents=True,
         verbose=False,
     )
+
     progress(0.9, desc="Done!")
     return qa_chain
 
@@ -148,10 +131,10 @@ def initialize_database(
     collection_name = create_collection_name(list_file_path[0])
 
     progress(0.25, desc="Loading document...")
-    doc_splits = load_doc(list_file_path, chunk_size, chunk_overlap)
+    doc_splits = load_doc_and_create_splits(list_file_path, chunk_size, chunk_overlap)
 
     progress(0.5, desc="Generating vector database...")
-    vector_db = create_db(doc_splits, collection_name)
+    vector_db = create_vector_db(doc_splits, collection_name)
 
     progress(0.9, desc="Done!")
     return vector_db, collection_name, "Complete!"