Create download_model.py

Add download_model.py script for pre-downloading Kimi-K2-Instruct model during build

download_model.py

@@ -0,0 +1,97 @@
+#!/usr/bin/env python3
+"""
+Model Downloader Script for Kimi-K2-Instruct
+This script pre-downloads the Kimi-K2-Instruct model from Hugging Face
+during the container build to speed up initialization.
+"""""
+
+import os
+import sys
+import logging
+from pathlib import Path
+
+# Setup logging
+logging.basicConfig(
+        level=logging.INFO,
+        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger(__name__)
+
+def download_model():
+        """Download the Kimi-K2-Instruct model from Hugging Face Hub"""""
+
+        try:
+                    from transformers import AutoTokenizer, AutoModelForCausalLM
+                    import torch
+
+                    MODEL_NAME = "moonshotai/Kimi-K2-Instruct"
+
+                    logger.info("=" * 60)
+                    logger.info("🤖 Starting Model Download")
+                    logger.info("=" * 60)
+
+                    # Check CUDA availability
+                    logger.info(f"CUDA Available: {torch.cuda.is_available()}")
+                    if torch.cuda.is_available():
+                                    logger.info(f"GPU Count: {torch.cuda.device_count()}")
+                                    for i in range(torch.cuda.device_count()):
+                                                        logger.info(f"  GPU {i}: {torch.cuda.get_device_name(i)}")
+
+                                                logger.info(f"\n📥 Downloading Model: {MODEL_NAME}")
+                                logger.info(f"Using HF_TOKEN: {'Yes' if os.environ.get('HF_TOKEN') else 'No'}")
+
+                    # Download tokenizer
+                    logger.info("\n[1/2] Downloading Tokenizer...")
+                    try:
+                                    tokenizer = AutoTokenizer.from_pretrained(
+                                                        MODEL_NAME,
+                                                        trust_remote_code=True,
+                                                        token=os.environ.get("HF_TOKEN")
+                                    )
+                                    logger.info("✓ Tokenizer downloaded successfully")
+                    except Exception as e:
+                                    logger.error(f"✗ Failed to download tokenizer: {str(e)}")
+                                    raise
+
+                                # Download model
+                                logger.info("\n[2/2] Downloading Model Weights (this may take a while)...")
+                    try:
+                                    model = AutoModelForCausalLM.from_pretrained(
+                                                        MODEL_NAME,
+                                                        torch_dtype=torch.bfloat16,
+                                                        device_map="auto",
+                                                        trust_remote_code=True,
+                                                        token=os.environ.get("HF_TOKEN")
+                                    )
+                                    logger.info("✓ Model downloaded successfully")
+                    except Exception as e:
+                                    logger.error(f"✗ Failed to download model: {str(e)}")
+                                    raise
+
+                                # Verify model is cached
+                                cache_dir = Path.home() / ".cache" / "huggingface" / "hub"
+                    logger.info(f"\n📁 Cache Directory: {cache_dir}")
+                    logger.info(f"✓ Model cached at: {cache_dir}")
+
+                    logger.info("\n" + "=" * 60)
+                    logger.info("✓ Model download completed successfully!")
+                    logger.info("=" * 60)
+
+                    return True
+
+        except ImportError as e:
+                    logger.error(f"✗ Import Error: {str(e)}")
+                    logger.error("Make sure all required packages are installed")
+                    return False
+        except Exception as e:
+                    logger.error(f"✗ Error: {str(e)}")
+                    import traceback
+                    traceback.print_exc()
+                    return False
+
+            if __name__ == "__main__":
+                    success = download_model()
+                    sys.exit(0 if success else 1)
+                                    )
+                                    )
+)