Update app.py

app.py

@@ -1,23 +1,42 @@
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 import gradio as gr
 import datetime
 import pytz
 import logging
 import gc
+import psutil
+import os
+from huggingface_hub import login
+import os
 
 # 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 = os.getenv('HF_TOKEN')
+    if not token:
+        # If not in environment, try to read from a file
+        token_file = 'hf_token.txt'
+        if os.path.exists(token_file):
+            with open(token_file, 'r') as f:
+                token = f.read().strip()
+
 class LocalLLMHandler:
     def __init__(self):
         self.model = None
         self.tokenizer = None
+        self.memory_tracker = MemoryTracker()
         
     def load_model(self, model_name="meta-llama/Llama-3.1-8B-Instruct"):
-        """Load model with CPU optimizations"""
+        """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