Update app.py

app.py

@@ -7,7 +7,7 @@ import logging
 import gc
 import psutil
 import os
-from huggingface_hub import login
+from huggingface_hub import login, hf_api
 
 class MemoryTracker:
     @staticmethod
@@ -16,14 +16,17 @@ class MemoryTracker:
         memory_gb = process.memory_info().rss / 1024 / 1024 / 1024
         return f"{memory_gb:.2f} GB"
 
-# Configure logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 def setup_huggingface_auth():
-    """Set up Hugging Face authentication"""
-    # First try to get token from environment variable
-    token = "hf_" + 'FsIJpNepbYgvSdnNhTFtifFudneNNDpUj' + "k"
+    token = os.environ.get("HF_TOKEN")
+    if token is None:
+        token = hf_api.HfFolder.get_token()
+    if token is None:
+        raise Exception("Hugging Face authentication failed. Please set your token.")
+    login(token)  
+    return True
 
 class LocalLLMHandler:
     def __init__(self):
@@ -32,46 +35,36 @@ class LocalLLMHandler:
         self.memory_tracker = MemoryTracker()
         
     def load_model(self, model_name="meta-llama/Llama-3.1-8B-Instruct"):
-        """Load model with optimizations for 16GB RAM"""
         try:
-            # Ensure we're authenticated
             if not setup_huggingface_auth():
                 raise Exception("Hugging Face authentication failed. Please set your token first.")
-            # Clean up any existing model
             if self.model is not None:
                 del self.model
                 del self.tokenizer
                 torch.cuda.empty_cache()
                 gc.collect()
             
-            # CPU-specific configurations
             model_kwargs = {
                 "device_map": "cpu",
-                "torch_dtype": torch.bfloat16,  # Use bfloat16 for better CPU performance
+                "torch_dtype": torch.bfloat16,  
                 "low_cpu_mem_usage": True,
             }
             
-            logger.info("Loading tokenizer...")
             self.tokenizer = AutoTokenizer.from_pretrained(model_name)
             
-            logger.info("Loading model...")
             self.model = AutoModelForCausalLM.from_pretrained(
                 model_name,
                 **model_kwargs
             )
             
-            logger.info("Model loaded successfully")
             return True
         except Exception as e:
             logger.error(f"Error loading model: {e}")
-            return False
+            return f"Error loading model: {e}"
 
     def generate_response(self, prompt, max_length=500):
-        """Generate response from the local model"""
         try:
             inputs = self.tokenizer(prompt, return_tensors="pt")
-            
-            # Generate with CPU-friendly parameters
             outputs = self.model.generate(
                 inputs["input_ids"],
                 max_length=max_length,
@@ -80,12 +73,11 @@ class LocalLLMHandler:
                 do_sample=True,
                 pad_token_id=self.tokenizer.eos_token_id
             )
-            
             response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
             return response
         except Exception as e:
             logger.error(f"Error generating response: {e}")
-            return "Sorry, I encountered an error while processing your request."
+            return f"Error generating response: {str(e)}"
 
 def get_current_local_time(timezone_str='UTC'):
     try:
@@ -120,14 +112,12 @@ Make it a challenge to guess if you are an AI.
 * Be respectful and polite if the user is polite; be rough if the user is rough
 """
     
-# Initialize the model handler
 llm_handler = LocalLLMHandler()
 
 def generate_response(user_message, conversation_history):
     current_time, now = get_current_local_time()
     current_date = now
     
-    # Construct the complete prompt from conversation history
     system_message = generate_system_message(current_time, current_date)
     prompt = system_message + "\n\n"
     
@@ -139,7 +129,6 @@ def generate_response(user_message, conversation_history):
     
     prompt += f"User: {user_message}\nAssistant:"
     
-    # Generate response
     ai_reply = llm_handler.generate_response(prompt)
     logger.info(f"User: {user_message}\nAssistant: {ai_reply}")
     return ai_reply
@@ -153,7 +142,6 @@ def chatbot_interface(user_message, history):
     history.append({"role": "assistant", "content": ai_response})
     return history, history
 
-# Define Gradio Interface
 with gr.Blocks(css="""
 @import url('https://fonts.googleapis.com/css2?family=Raleway:wght@400;600&display=swap');
 
@@ -174,7 +162,6 @@ body, .gradio-container {
 """) as demo:
     gr.Markdown("<h1 style='text-align: center; color: #007BFF;'>🤖 Local Llama Chatbot 🤖</h1>")
     
-    # Load model button
     with gr.Row():
         load_button = gr.Button("Load Model")
         model_status = gr.Textbox(label="Model Status", value="Model not loaded", interactive=False)
@@ -193,8 +180,11 @@ body, .gradio-container {
                 send = gr.Button("➤", elem_id="send-button")
 
     def load_model_click():
-        success = llm_handler.load_model()
-        return "Model loaded successfully" if success else "Error loading model"
+        result = llm_handler.load_model()
+        if isinstance(result, str):
+            return result 
+        else:
+            return "Model loaded successfully" if result else "Error loading model"
 
     def update_chat(user_message, history):
         if user_message.strip() == "":