import os
import logging
from transformers import AutoTokenizer, AutoModelForCausalLM
from huggingface_hub import snapshot_download

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def download_phi2():
    """Scarica Microsoft Phi-2 (2.7B) - Ottimo bilanciamento qualità/dimensione"""
    model_name = "microsoft/phi-2"
    local_path = "./models/phi-2"
    
    try:
        os.makedirs(local_path, exist_ok=True)
        
        logger.info("📥 Scaricando Microsoft Phi-2 (2.7B parametri)...")
        logger.info("💾 Dimensione stimata: ~5GB")
        
        snapshot_download(
            repo_id=model_name,
            local_dir=local_path,
            local_dir_use_symlinks=False,
            resume_download=True
        )
        
        logger.info("✅ Phi-2 scaricato correttamente!")
        
        # Test
        tokenizer = AutoTokenizer.from_pretrained(local_path, local_files_only=True, trust_remote_code=True)
        logger.info(f"✅ Tokenizer pronto: {len(tokenizer)} tokens")
        
        return local_path
        
    except Exception as e:
        logger.error(f"Errore: {e}")
        return None

if __name__ == "__main__":
    print("=== DOWNLOAD PHI-2 ===")
    path = download_phi2()
    if path:
        print(f"\nUsa: python main.py --mode worker --model-path '{path}' --gpu-id 0")