# download_models.py
from huggingface_hub import snapshot_download
import os
import shutil  # [新增] 引入 shutil 模块用于删除非空目录
import time


def download_model(model_name, local_dir):
    """
    下载Hugging Face模型到本地目录

    Args:
        model_name: Hugging Face模型名称
        local_dir: 本地保存目录
    """
    print(f"开始下载模型: {model_name}")
    print(f"保存到: {local_dir}")

    try:
        # 创建目录
        if os.path.exists(local_dir):
            # [修改] 使用 shutil.rmtree 删除目录及其内部所有内容
            shutil.rmtree(local_dir)
            print(f"目录 {local_dir} 已删除")
            time.sleep(5)
        os.makedirs(local_dir, exist_ok=True)

        # 下载模型
        snapshot_download(
            repo_id=model_name,
            local_dir=local_dir,
            local_dir_use_symlinks=False,  # 不使用符号链接
            ignore_patterns=["*.safetensors"],  # 可选：不下载safetensors格式
        )

        print(f"✓ 模型 {model_name} 下载完成")
        print(f"保存路径: {os.path.abspath(local_dir)}")

    except Exception as e:
        print(f"✗ 下载失败: {e}")


if __name__ == "__main__":
    # 定义要下载的模型
    models_to_download = [
        ("bert-base-uncased", "./models/bert-base-uncased"),
        ("bert-base-chinese", "./models/bert-base-chinese"),
    ]

    print("开始下载BERT模型...")
    print("=" * 50)

    for model_name, local_dir in models_to_download:
        download_model(model_name, local_dir)
        print()

    print("=" * 50)
    print("所有模型下载完成！")
    print("\n模型文件结构:")
    for model_name, local_dir in models_to_download:
        if os.path.exists(local_dir):
            file_count = len([f for f in os.listdir(local_dir) if os.path.isfile(os.path.join(local_dir, f))])
            print(f"  {local_dir}: {file_count} 个文件")

    print("\n使用示例:")
    print("本地模型路径:")
    print("  - 英文: './models/bert-base-uncased'")
    print("  - 中文: './models/bert-base-chinese'")