Create app/setup_models.py

app/setup_models.py

@@ -0,0 +1,44 @@
+import logging
+import os
+
+# 1. Use the Mirror to bypass connection issues
+os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
+# 2. Increase timeouts so it doesn't fail on slow handshakes
+os.environ["HF_HUB_DOWNLOAD_TIMEOUT"] = "300"
+os.environ["HF_HUB_ETAG_TIMEOUT"] = "300"
+
+from transformers import AutoModel, AutoModelForCausalLM, AutoTokenizer
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger("ModelDownloader")
+
+
+def download_everything():
+    logger.info("🚀 STARTING ROBUST DOWNLOAD...")
+
+    # 1. Download CodeBERT
+    logger.info("📦 Downloading CodeBERT (Microsoft)...")
+    try:
+        AutoTokenizer.from_pretrained("microsoft/codebert-base")
+        AutoModel.from_pretrained("microsoft/codebert-base")
+        logger.info("✅ CodeBERT Cached!")
+    except Exception as e:
+        logger.error(f"❌ CodeBERT Failed: {e}")
+
+    # 2. Download Qwen
+    logger.info("🧠 Downloading Qwen 1.5B (This is ~3GB, please wait)...")
+    try:
+        model_id = "Qwen/Qwen2.5-Coder-1.5B-Instruct"
+        AutoTokenizer.from_pretrained(model_id)
+        # This is the heavy part that was failing
+        AutoModelForCausalLM.from_pretrained(model_id)
+        logger.info("✅ Qwen 1.5B Cached!")
+    except Exception as e:
+        logger.error(f"❌ Qwen Failed: {e}")
+        logger.info(
+            "💡 TIP: If it fails, just run this script again. It will resume where it left off."
+        )
+
+
+if __name__ == "__main__":
+    download_everything()