move_models.py

#!/usr/bin/env python3
"""
模型迁移工具
将已下载的 HuggingFace 模型从默认缓存目录移动到项目的 models 目录
"""

import os
import shutil
from pathlib import Path


def find_huggingface_cache():
    """查找 HuggingFace 默认缓存目录"""
    possible_paths = [
        Path.home() / ".cache" / "huggingface",
        Path.home() / ".cache" / "huggingface" / "hub",
    ]
    
    for path in possible_paths:
        if path.exists():
            return path
    
    return None


def list_cached_models(cache_dir):
    """列出缓存目录中的所有模型"""
    hub_dir = cache_dir / "hub" if (cache_dir / "hub").exists() else cache_dir
    
    models = []
    if hub_dir.exists():
        for item in hub_dir.iterdir():
            if item.is_dir() and item.name.startswith("models--"):
                # 从目录名提取模型名称
                # 例如: models--Qwen--Qwen2.5-1.5B-Instruct -> Qwen/Qwen2.5-1.5B-Instruct
                model_name = item.name.replace("models--", "").replace("--", "/")
                
                # 计算模型大小
                size = sum(f.stat().st_size for f in item.rglob("*") if f.is_file())
                size_mb = size / (1024 * 1024)
                
                models.append({
                    "name": model_name,
                    "path": item,
                    "size_mb": size_mb
                })
    
    return models


def move_model(source_path, target_dir, model_name):
    """移动模型到目标目录"""
    # HuggingFace 缓存需要 hub 子目录
    # 创建目标路径: models/hub/models--Qwen--Qwen2.5-1.5B-Instruct
    hub_dir = target_dir / "hub"
    hub_dir.mkdir(exist_ok=True)
    
    target_path = hub_dir / source_path.name
    
    if target_path.exists():
        print(f"  ⚠️  目标已存在: {target_path}")
        response = input("  是否覆盖? (y/n): ")
        if response.lower() != 'y':
            print("  跳过")
            return False
        shutil.rmtree(target_path)
    
    print(f"  正在移动: {source_path}")
    print(f"  目标位置: {target_path}")
    
    try:
        shutil.move(str(source_path), str(target_path))
        print("  ✅ 移动成功!")
        return True
    except Exception as e:
        print(f"  ❌ 移动失败: {e}")
        return False


def main():
    print("=" * 70)
    print("🔄 HuggingFace 模型迁移工具")
    print("=" * 70)
    print()
    
    # 查找缓存目录
    cache_dir = find_huggingface_cache()
    if not cache_dir:
        print("❌ 未找到 HuggingFace 缓存目录")
        print("通常位于: ~/.cache/huggingface/")
        return
    
    print(f"📂 HuggingFace 缓存目录: {cache_dir}")
    print()
    
    # 列出已缓存的模型
    models = list_cached_models(cache_dir)
    
    if not models:
        print("ℹ️  缓存目录中没有找到模型")
        return
    
    print(f"找到 {len(models)} 个已下载的模型:")
    print()
    
    for i, model in enumerate(models, 1):
        print(f"{i}. {model['name']}")
        print(f"   大小: {model['size_mb']:.2f} MB")
        print(f"   路径: {model['path']}")
        print()
    
    # 设置目标目录
    project_root = Path(__file__).parent
    target_dir = project_root / "models"
    target_dir.mkdir(exist_ok=True)
    
    print(f"📁 目标目录: {target_dir}")
    print()
    
    # 询问用户要移动哪些模型
    print("选择要移动的模型:")
    print("  - 输入模型编号 (例如: 1)")
    print("  - 输入 'all' 移动所有模型")
    print("  - 输入 'q' 退出")
    print()
    
    choice = input("请选择: ").strip().lower()
    
    if choice == 'q':
        print("👋 已取消")
        return
    
    models_to_move = []
    
    if choice == 'all':
        models_to_move = models
    else:
        try:
            index = int(choice) - 1
            if 0 <= index < len(models):
                models_to_move = [models[index]]
            else:
                print("❌ 无效的编号")
                return
        except ValueError:
            print("❌ 无效的输入")
            return
    
    print()
    print(f"准备移动 {len(models_to_move)} 个模型...")
    print()
    
    # 移动模型
    success_count = 0
    for model in models_to_move:
        print(f"处理: {model['name']}")
        if move_model(model['path'], target_dir, model['name']):
            success_count += 1
        print()
    
    print("=" * 70)
    print(f"✅ 完成! 成功移动 {success_count}/{len(models_to_move)} 个模型")
    print("=" * 70)
    print()
    print("提示:")
    print("  1. 模型已移动到项目的 models 目录")
    print("  2. 下次运行程序时将直接使用这些模型")
    print("  3. 可以手动删除原缓存目录以释放空间")
    print(f"     rm -rf {cache_dir / 'hub'}")
    print()


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print("\n\n👋 已取消")
    except Exception as e:
        print(f"\n❌ 发生错误: {e}")
        import traceback
        traceback.print_exc()