Fix offline mode and improve model loading

Dockerfile

@@ -26,21 +26,17 @@ ENV PYTHONDONTWRITEBYTECODE=1
 ENV PORT=7860
 ENV PATH="/home/user/.local/bin:${PATH}"
 ENV HF_HOME=/home/user/.cache/huggingface
-ENV TRANSFORMERS_CACHE=/home/user/.cache/huggingface
-# Set logging to stdout
-ENV LOG_FILE=/dev/stdout
 # Memory optimizations
 ENV MALLOC_ARENA_MAX=2
 ENV MALLOC_TRIM_THRESHOLD_=100000
 ENV MALLOC_MMAP_THRESHOLD_=100000
-# Transformers optimizations
-ENV TRANSFORMERS_OFFLINE=1
-ENV TORCH_CUDA_ARCH_LIST="3.5;5.0;6.0;6.1;7.0;7.5;8.0;8.6+PTX"
-ENV CUDA_LAUNCH_BLOCKING=1
 # Model optimizations
 ENV OMP_NUM_THREADS=1
 ENV MKL_NUM_THREADS=1
 ENV NUMEXPR_NUM_THREADS=1
+# Ensure offline mode is disabled
+ENV HF_HUB_OFFLINE=0
+ENV TRANSFORMERS_OFFLINE=0
 
 # Switch to non-root user
 USER user

src/model_manager.py

@@ -3,6 +3,7 @@ from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 import torch
 from huggingface_hub import login
 from .config import Config
+import os
 
 logger = logging.getLogger(__name__)
 
@@ -13,11 +14,15 @@ class ModelManager:
         self.model = None
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
         
+        # Ensure offline mode is disabled
+        os.environ['HF_HUB_OFFLINE'] = '0'
+        os.environ['TRANSFORMERS_OFFLINE'] = '0'
+        
         # Login to Hugging Face Hub
         if Config.HUGGING_FACE_TOKEN:
             logger.info("Logging in to Hugging Face Hub")
             try:
-                login(token=Config.HUGGING_FACE_TOKEN)
+                login(token=Config.HUGGING_FACE_TOKEN, add_to_git_credential=False)
                 logger.info("Successfully logged in to Hugging Face Hub")
             except Exception as e:
                 logger.error(f"Failed to login to Hugging Face Hub: {str(e)}")
@@ -34,7 +39,8 @@ class ModelManager:
             self.tokenizer = AutoTokenizer.from_pretrained(
                 self.model_name,
                 token=Config.HUGGING_FACE_TOKEN,
-                model_max_length=1024  # Limit max length to save memory
+                model_max_length=1024,  # Limit max length to save memory
+                trust_remote_code=True
             )
             # Ensure we have the necessary special tokens
             special_tokens = {
@@ -71,14 +77,8 @@ class ModelManager:
                 token=Config.HUGGING_FACE_TOKEN,
                 low_cpu_mem_usage=True,
                 torch_dtype=torch.float16,  # Use fp16 for additional memory savings
-                max_memory={0: "4GB"},  # Limit memory usage
-                offload_folder="offload",  # Enable CPU offloading
-                use_cache=False  # Disable KV cache to save memory
+                trust_remote_code=True
             )
-
-            # Enable gradient checkpointing
-            self.model.gradient_checkpointing_enable()
-            
             # Resize embeddings to match tokenizer
             self.model.resize_token_embeddings(len(self.tokenizer))
             logger.info("Model loaded successfully")