up

chatbot.py

@@ -318,27 +318,26 @@ Response:"""
         
     def _load_emotion_model(self):
         try:
-            # Use safetensors explicitly for the primary model
+            # Load emotion model directly from Hugging Face
             return pipeline(
                 "text-classification",
                 model="SamLowe/roberta-base-go_emotions",
                 top_k=None,
-                device_map="auto" if self.device == "cuda" else None,
+                device_map="auto" if torch.cuda.is_available() else None,
                 cache_dir=CACHE_DIR,
-                use_safetensors=True  # Explicitly use safetensors
+                local_files_only=False  # Ensure we download from Hugging Face
             )
         except Exception as e:
             logger.error(f"Error loading emotion model: {e}")
-            # Fallback to a simpler model with safetensors
+            # Fallback to a simpler model
             try:
                 return pipeline(
                     "text-classification", 
                     model="j-hartmann/emotion-english-distilroberta-base", 
                     return_all_scores=True,
-                    device_map="auto" if self.device == "cuda" else None,
+                    device_map="auto" if torch.cuda.is_available() else None,
                     cache_dir=CACHE_DIR,
-                    use_safetensors=True,  # Explicitly use safetensors
-                    from_tf=True  # Use TensorFlow weights if available
+                    local_files_only=False  # Ensure we download from Hugging Face
                 )
             except Exception as e:
                 logger.error(f"Error loading fallback emotion model: {e}")
@@ -347,8 +346,8 @@ Response:"""
     
     def _initialize_llm(self, model_name: str, use_4bit: bool):
         try:
-            # Configure quantization if needed
-            if use_4bit:
+            # Configure quantization only if CUDA is available
+            if use_4bit and torch.cuda.is_available():
                 quantization_config = BitsAndBytesConfig(
                     load_in_4bit=True,
                     bnb_4bit_compute_dtype=torch.float16,
@@ -357,22 +356,25 @@ Response:"""
                 )
             else:
                 quantization_config = None
+                logger.info("CUDA not available, running without quantization")
 
             # Load base model
             logger.info(f"Loading base model: {model_name}")
             base_model = AutoModelForCausalLM.from_pretrained(
                 model_name,
                 quantization_config=quantization_config,
-                device_map="auto" if self.device == "cuda" else None,
+                device_map="auto" if torch.cuda.is_available() else None,
                 trust_remote_code=True,
-                cache_dir=CACHE_DIR
+                cache_dir=CACHE_DIR,
+                use_auth_token=os.environ.get('HF_TOKEN')  # Add auth token for gated models
             )
 
             # Load tokenizer
             logger.info("Loading tokenizer")
             tokenizer = AutoTokenizer.from_pretrained(
                 model_name,
-                cache_dir=CACHE_DIR
+                cache_dir=CACHE_DIR,
+                use_auth_token=os.environ.get('HF_TOKEN')  # Add auth token for gated models
             )
             tokenizer.pad_token = tokenizer.eos_token
 
@@ -381,7 +383,8 @@ Response:"""
             model = PeftModel.from_pretrained(
                 base_model, 
                 self.peft_model_path,
-                cache_dir=CACHE_DIR
+                cache_dir=CACHE_DIR,
+                use_auth_token=os.environ.get('HF_TOKEN')  # Add auth token for gated models
             )
             logger.info("Successfully loaded PEFT model")
 
@@ -395,7 +398,7 @@ Response:"""
                 top_p=0.95,
                 repetition_penalty=1.1,
                 do_sample=True,
-                device_map="auto" if self.device == "cuda" else None
+                device_map="auto" if torch.cuda.is_available() else None
             )
             
             # Create LangChain wrapper