Fix indentation issues in download_model.py

download_model.py

@@ -1,90 +1,58 @@
 #!/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.
+This script pre-downloads the model from Hugging Face.
 """""
 
 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)
+                    print("=" * 60)
+                    print("Starting Model Download")
+                    print("=" * 60)
 
-                    # Check CUDA availability
-                    logger.info(f"CUDA Available: {torch.cuda.is_available()}")
+                    print(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)}")
+                                    print(f"GPU Count: {torch.cuda.device_count()}")
 
-                                                logger.info(f"\n📥 Downloading Model: {MODEL_NAME}")
-                                logger.info(f"Using HF_TOKEN: {'Yes' if os.environ.get('HF_TOKEN') else 'No'}")
+                    print(f"Downloading Model: {MODEL_NAME}")
 
                     # 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
+                    print("Downloading Tokenizer...")
+                    tokenizer = AutoTokenizer.from_pretrained(
+                                    MODEL_NAME,
+                                    trust_remote_code=True,
+                                    token=os.environ.get("HF_TOKEN")
+                    )
+                    print("Tokenizer downloaded successfully")
 
-                                # 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
+                    # Download model
+                    print("Downloading Model Weights...")
+                    model = AutoModelForCausalLM.from_pretrained(
+                                    MODEL_NAME,
+                                    torch_dtype=torch.bfloat16,
+                                    device_map="auto",
+                                    trust_remote_code=True,
+                                    token=os.environ.get("HF_TOKEN")
+                    )
+                    print("Model downloaded successfully")
 
-                                # 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)
+                    print("=" * 60)
+                    print("Model download completed!")
+                    print("=" * 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)}")
+                    print(f"Error: {str(e)}")
                     import traceback
                     traceback.print_exc()
                     return False
@@ -92,6 +60,5 @@ def download_model():
             if __name__ == "__main__":
                     success = download_model()
                     sys.exit(0 if success else 1)
-                                    )
-                                    )
-)
+                    )
+                    )