up

chatbot.py

@@ -6,6 +6,7 @@ import re
 from typing import List, Dict, Any, Optional, Union
 from datetime import datetime
 from pydantic import BaseModel, Field
+import tempfile
 
 # Model imports
 from transformers import (
@@ -34,7 +35,7 @@ from conversation_flow import FlowManager
 logging.basicConfig(
     level=logging.INFO,
     format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
-    handlers=[logging.StreamHandler()]  # Only use StreamHandler for Hugging Face Spaces
+    handlers=[logging.StreamHandler()]
 )
 logger = logging.getLogger(__name__)
 
@@ -42,12 +43,36 @@ logger = logging.getLogger(__name__)
 import warnings
 warnings.filterwarnings('ignore', category=UserWarning)
 
-# Set environment variables
-os.environ.update({
-    'TRANSFORMERS_VERBOSITY': 'error',
-    'TOKENIZERS_PARALLELISM': 'false',
-    'BITSANDBYTES_NOWELCOME': '1'
-})
+# Set up cache directories
+def setup_cache_dirs():
+    # Check if running in Hugging Face Spaces
+    is_spaces = os.environ.get('SPACE_ID') is not None
+    
+    if is_spaces:
+        # Use /tmp for Hugging Face Spaces
+        cache_dir = '/tmp/huggingface'
+        os.environ.update({
+            'TRANSFORMERS_CACHE': cache_dir,
+            'HF_HOME': cache_dir,
+            'TOKENIZERS_PARALLELISM': 'false',
+            'TRANSFORMERS_VERBOSITY': 'error',
+            'BITSANDBYTES_NOWELCOME': '1'
+        })
+    else:
+        # Use default cache for local development
+        cache_dir = os.path.expanduser('~/.cache/huggingface')
+        os.environ.update({
+            'TOKENIZERS_PARALLELISM': 'false',
+            'TRANSFORMERS_VERBOSITY': 'error',
+            'BITSANDBYTES_NOWELCOME': '1'
+        })
+    
+    # Create cache directory if it doesn't exist
+    os.makedirs(cache_dir, exist_ok=True)
+    return cache_dir
+
+# Set up cache directories
+CACHE_DIR = setup_cache_dirs()
 
 # Define base directory and paths
 BASE_DIR = os.path.abspath(os.path.dirname(__file__))
@@ -289,19 +314,17 @@ Response:"""
                 model="SamLowe/roberta-base-go_emotions",
                 top_k=None,
                 device_map="auto" if self.device == "cuda" else None,
-                local_files_only=False,  # Force download from Hugging Face Hub
-                use_auth_token=False  # No authentication needed for public models
+                cache_dir=CACHE_DIR
             )
         except Exception as e:
             logger.error(f"Error loading emotion model: {e}")
-            # Fallback to a simpler model
+            # Fallback 
             return pipeline(
                 "text-classification", 
                 model="j-hartmann/emotion-english-distilroberta-base", 
                 return_all_scores=True,
                 device_map="auto" if self.device == "cuda" else None,
-                local_files_only=False,
-                use_auth_token=False
+                cache_dir=CACHE_DIR
             )
     
     def _initialize_llm(self, model_name: str, use_4bit: bool):
@@ -317,33 +340,30 @@ Response:"""
             else:
                 quantization_config = None
 
-            # Load base model directly from Hugging Face Hub
+            # 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,
                 trust_remote_code=True,
-                local_files_only=False,
-                use_auth_token=False
+                cache_dir=CACHE_DIR
             )
 
-            # Load tokenizer directly from Hugging Face Hub
+            # Load tokenizer
             logger.info("Loading tokenizer")
             tokenizer = AutoTokenizer.from_pretrained(
                 model_name,
-                local_files_only=False,
-                use_auth_token=False
+                cache_dir=CACHE_DIR
             )
             tokenizer.pad_token = tokenizer.eos_token
 
-            # Load PEFT model directly from Hugging Face Hub
+            # Load PEFT model
             logger.info(f"Loading PEFT model from {self.peft_model_path}")
             model = PeftModel.from_pretrained(
                 base_model, 
                 self.peft_model_path,
-                local_files_only=False,
-                use_auth_token=False
+                cache_dir=CACHE_DIR
             )
             logger.info("Successfully loaded PEFT model")
 

requirements.txt

@@ -24,4 +24,5 @@ tokenizers>=0.21.1
 tiktoken>=0.9.0
 starlette>=0.46.1
 websockets>=15.0.1
+tempfile>=0.1.0
 python-multipart>=0.0.6