Update Dockerfile

app.py

@@ -128,33 +128,29 @@ CACHE_DIR = "/cache/models"
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
 def load_model():
-    """Load model directly, attempting cache first"""
+    """Load model with automatic cache handling"""
     try:
-        # Try loading from cache
-        print("Attempting to load from cache...")
-        model = AutoModelForCausalLM.from_pretrained(
+        # First try with local files only (uses cache if available)
+        print("Checking for cached model...")
+        return AutoModelForCausalLM.from_pretrained(
             MODEL_ID,
             cache_dir=CACHE_DIR,
-            local_files_only=True  # Force cache usage
-        ).to(DEVICE)
-        tokenizer = AutoTokenizer.from_pretrained(
+            local_files_only=True  # Will fail if not cached
+        ).to(DEVICE), AutoTokenizer.from_pretrained(
             MODEL_ID,
             cache_dir=CACHE_DIR,
             local_files_only=True
         )
     except OSError:
-        # Fallback to download if cache missing
-        print("Cache not found, downloading...")
-        model = AutoModelForCausalLM.from_pretrained(
+        # Fallback to download if not in cache
+        print("Downloading model...")
+        return AutoModelForCausalLM.from_pretrained(
             MODEL_ID,
             cache_dir=CACHE_DIR
-        ).to(DEVICE)
-        tokenizer = AutoTokenizer.from_pretrained(
+        ).to(DEVICE), AutoTokenizer.from_pretrained(
             MODEL_ID,
             cache_dir=CACHE_DIR
         )
-    
-    return model, tokenizer
 
 # Load model
 model, tokenizer = load_model()