Update Dockerfile

app.py

@@ -134,34 +134,39 @@ def load_model_with_retry(max_retries=3):
             # First try loading from local cache only
             if attempt == 0:
                 print("Attempting to load from local cache...")
-                model = AutoModelForCausalLM.from_pretrained(
-                    MODEL_ID,
-                    cache_dir=CACHE_DIR,
-                    local_files_only=True  # Only use cache
-                ).to(DEVICE)
-                tokenizer = AutoTokenizer.from_pretrained(
-                    MODEL_ID,
-                    cache_dir=CACHE_DIR,
-                    local_files_only=True
-                )
-            # Fallback to download if cache missing
-            else:
-                print(f"Downloading model (attempt {attempt + 1})...")
-                model = AutoModelForCausalLM.from_pretrained(
-                    MODEL_ID,
-                    cache_dir=CACHE_DIR
-                ).to(DEVICE)
-                tokenizer = AutoTokenizer.from_pretrained(
-                    MODEL_ID,
-                    cache_dir=CACHE_DIR
-                )
+                try:
+                    model = AutoModelForCausalLM.from_pretrained(
+                        MODEL_ID,
+                        cache_dir=CACHE_DIR,
+                        local_files_only=True
+                    ).to(DEVICE)
+                    tokenizer = AutoTokenizer.from_pretrained(
+                        MODEL_ID,
+                        cache_dir=CACHE_DIR,
+                        local_files_only=True
+                    )
+                    return model, tokenizer
+                except OSError:
+                    print("Cache not found, will download...")
+                    continue
+
+            # Download with retry
+            print(f"Downloading model (attempt {attempt + 1})...")
+            model = AutoModelForCausalLM.from_pretrained(
+                MODEL_ID,
+                cache_dir=CACHE_DIR
+            ).to(DEVICE)
+            tokenizer = AutoTokenizer.from_pretrained(
+                MODEL_ID,
+                cache_dir=CACHE_DIR
+            )
             return model, tokenizer
             
         except Exception as e:
             if attempt == max_retries - 1:
                 raise RuntimeError(f"Failed after {max_retries} attempts: {str(e)}")
             
-            wait_time = min(2 ** (attempt + 1), 10)  # Cap at 10 seconds
+            wait_time = min(2 ** (attempt + 1), 10)
             print(f"Attempt {attempt + 1} failed ({str(e)}), retrying in {wait_time}s...")
             time.sleep(wait_time)