download_models.py

#!/usr/bin/env python3
"""Download models at Docker build time"""

import os
import requests
from pathlib import Path

MODELS_DIR = "models"
os.makedirs(MODELS_DIR, exist_ok=True)

MODEL_CONFIGS = {
    "fast-chat": {
        "file": "qwen2.5-0.5b-instruct-q4_k_m.gguf",
        "url": "https://huggingface.co/Qwen/Qwen2.5-0.5B-Instruct-GGUF/resolve/main/qwen2.5-0.5b-instruct-q4_k_m.gguf"
    },
    "tinyllama": {
        "file": "tinyllama-1.1b-chat-v1.0.Q4_K_M.gguf",
        "url": "https://huggingface.co/TheBloke/TinyLlama-1.1B-Chat-v1.0-GGUF/resolve/main/tinyllama-1.1b-chat-v1.0.Q4_K_M.gguf"
    },
    "coder": {
        "file": "qwen2.5-coder-1.5b-instruct-q4_k_m.gguf",
        "url": "https://huggingface.co/Qwen/Qwen2.5-Coder-1.5B-Instruct-GGUF/resolve/main/qwen2.5-coder-1.5b-instruct-q4_k_m.gguf"
    }
}

def download_model(model_id, config):
    """Download a single model"""
    filepath = os.path.join(MODELS_DIR, config["file"])
    
    # Skip if already exists and has reasonable size
    if os.path.exists(filepath) and os.path.getsize(filepath) > 50000000:
        print(f"✓ {model_id} already exists ({os.path.getsize(filepath) / 1e9:.2f}GB)")
        return
    
    print(f"Downloading {model_id}...")
    try:
        response = requests.get(config["url"], stream=True, timeout=60)
        response.raise_for_status()
        
        total_size = int(response.headers.get('content-length', 0))
        downloaded = 0
        
        with open(filepath, 'wb') as f:
            for chunk in response.iter_content(chunk_size=10*1024*1024):  # 10MB chunks
                if chunk:
                    f.write(chunk)
                    downloaded += len(chunk)
                    if total_size:
                        pct = (downloaded / total_size) * 100
                        print(f"  {model_id}: {pct:.1f}%", end='\r')
        
        print(f"✓ {model_id} downloaded ({os.path.getsize(filepath) / 1e9:.2f}GB)")
    except Exception as e:
        print(f"✗ Failed to download {model_id}: {e}")

if __name__ == "__main__":
    print("Pre-downloading models at build time...")
    
    # Only download fast-chat at build time (others on-demand)
    download_model("fast-chat", MODEL_CONFIGS["fast-chat"])
    
    print(f"\n✓ Models ready in {MODELS_DIR}/")