Update app.py

app.py

@@ -16,6 +16,7 @@ import chromadb
 import tempfile
 from tenacity import retry, stop_after_attempt, wait_exponential, retry_if_exception_type
 import requests
+from transformers import BitsAndBytesConfig
 
 # Set up logging
 logging.basicConfig(level=logging.INFO)
@@ -28,9 +29,9 @@ if os.environ["HUGGINGFACEHUB_API_TOKEN"] == "default-token":
 
 # Model and embedding options
 LLM_MODELS = {
-    "Balanced (Mixtral-8x7B)": "mistralai/Mixtral-8x7B-Instruct-v0.1",
-    "Lightweight (Gemma-2B)": "google/gemma-2-2b-it",
-    "High Accuracy (Llama-3-8B)": "meta-llama/Llama-3-8b-hf"
+    "High Accuracy (Mixtral-8x7B)": "mistralai/Mixtral-8x7B-Instruct-v0.1",
+    "Balanced (Gemma-2-9B)": "google/gemma-2-9b-it",
+    "Lightweight (Mistral-7B)": "mistralai/Mistral-7B-Instruct-v0.2"
 }
 
 EMBEDDING_MODELS = {
@@ -160,13 +161,16 @@ def initialize_qa_chain(llm_model, temperature):
         return "Please process documents first.", None
 
     try:
+        # Enable quantization for Mixtral-8x7B to reduce memory usage
+        quantization_config = BitsAndBytesConfig(load_in_4bit=True) if llm_model == "mistralai/Mixtral-8x7B-Instruct-v0.1" else None
         llm = HuggingFaceEndpoint(
             repo_id=LLM_MODELS[llm_model],
             task="text-generation",
             temperature=float(temperature),
             max_new_tokens=512,
             huggingfacehub_api_token=os.environ["HUGGINGFACEHUB_API_TOKEN"],
-            timeout=30
+            timeout=30,
+            quantization_config=quantization_config
         )
         # Dynamically set k based on vector store size
         collection = vector_store._collection
@@ -182,13 +186,13 @@ def initialize_qa_chain(llm_model, temperature):
     except requests.exceptions.HTTPError as e:
         logger.error(f"HTTP error initializing QA chain for {llm_model}: {str(e)}")
         if "503" in str(e):
-            return f"Error: Hugging Face API temporarily unavailable for {llm_model}. Try 'Balanced (Mixtral-8x7B)' or wait and retry.", None
+            return f"Error: Hugging Face API temporarily unavailable for {llm_model}. Try 'Lightweight (Mistral-7B)' or wait and retry.", None
         elif "403" in str(e):
-            return f"Error: Access denied for {llm_model}. Ensure your HF token has access.", None
+            return f"Error: Access denied for {llm_model}. Ensure your HF token is valid.", None
         return f"Error initializing QA chain: {str(e)}.", None
     except Exception as e:
         logger.error(f"Error initializing QA chain for {llm_model}: {str(e)}")
-        return f"Error initializing QA chain: {str(e)}. Ensure your HF token has access to {llm_model}.", None
+        return f"Error initializing QA chain: {str(e)}. Ensure your HF token is valid.", None
 
 # Function to handle user query with retry logic
 @retry(
@@ -214,9 +218,9 @@ def answer_question(question, llm_model, embedding_model, temperature, chunk_siz
     except requests.exceptions.HTTPError as e:
         logger.error(f"HTTP error answering question: {str(e)}")
         if "503" in str(e):
-            return f"Error: Hugging Face API temporarily unavailable for {llm_model}. Try 'Balanced (Mixtral-8x7B)' or wait and retry.", chat_history
+            return f"Error: Hugging Face API temporarily unavailable for {llm_model}. Try 'Lightweight (Mistral-7B)' or wait and retry.", chat_history
         elif "403" in str(e):
-            return f"Error: Access denied for {llm_model}. Ensure your HF token has access.", chat_history
+            return f"Error: Access denied for {llm_model}. Ensure your HF token is valid.", chat_history
         return f"Error answering question: {str(e)}", chat_history
     except Exception as e:
         logger.error(f"Error answering question: {str(e)}")
@@ -272,7 +276,7 @@ with gr.Blocks(theme=gr.themes.Soft(), title="DocTalk: Document Q&A Chatbot") as
             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="Balanced (Mixtral-8x7B)")
+            llm_model = gr.Dropdown(choices=list(LLM_MODELS.keys()), label="Select LLM Model", value="High Accuracy (Mixtral-8x7B)")
             embedding_model = gr.Dropdown(choices=list(EMBEDDING_MODELS.keys()), label="Select Embedding Model", value="Lightweight (MiniLM-L6)")
             temperature = gr.Slider(minimum=0.1, 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")