Direction9/model/verify_models.py

#!/usr/bin/env python3
"""
验证已下载的模型是否完整可用
支持文件完整性检查 + 实际推理测试
"""

import json
import fnmatch
from pathlib import Path

PROJECT_ROOT = Path(__file__).parent.parent.resolve()
WEIGHTS_DIR = PROJECT_ROOT / "model" / "weights"

# 必需的核心文件
REQUIRED_FILES = {
    "config.json",
    "tokenizer.json",
    "tokenizer_config.json",
}

# 模型权重文件（至少有一种）
WEIGHT_FILES = {
    "pytorch_model.bin",
    "model.safetensors",
    "tf_model.h5",
    "model.onnx",
}

# 分片权重文件模式
SHARDED_WEIGHT_PATTERNS = [
    "model-*.safetensors",
    "pytorch_model-*.bin",
    "model-*.bin",
]


def verify_model(model_dir: Path) -> dict:
    """验证单个模型目录"""
    result = {
        "name": model_dir.name,
        "path": str(model_dir),
        "exists": model_dir.exists(),
        "files": [],
        "missing_required": [],
        "has_weights": False,
        "weight_files": [],
        "valid": False,
        "error": None
    }

    if not model_dir.exists():
        result["error"] = "模型目录不存在"
        return result

    # 列出所有文件
    all_files = {f.name for f in model_dir.iterdir() if f.is_file()}
    result["files"] = sorted(all_files)

    # 检查必需文件
    for req_file in REQUIRED_FILES:
        if req_file not in all_files:
            found = False
            for sub_dir in model_dir.iterdir():
                if sub_dir.is_dir():
                    sub_files = {f.name for f in sub_dir.iterdir() if f.is_file()}
                    if req_file in sub_files:
                        found = True
                        break
            if not found:
                result["missing_required"].append(req_file)

    # 检查权重文件
    for weight_file in WEIGHT_FILES:
        if weight_file in all_files:
            result["weight_files"].append(weight_file)
            result["has_weights"] = True
        else:
            for sub_dir in model_dir.iterdir():
                if sub_dir.is_dir():
                    sub_files = {f.name for f in sub_dir.iterdir() if f.is_file()}
                    if weight_file in sub_files:
                        result["weight_files"].append(f"{sub_dir.name}/{weight_file}")
                        result["has_weights"] = True
                        break

    # 检查分片权重文件
    for pattern in SHARDED_WEIGHT_PATTERNS:
        for f in model_dir.iterdir():
            if f.is_file() and fnmatch.fnmatch(f.name, pattern):
                result["weight_files"].append(f.name)
                result["has_weights"] = True

    # 验证完整性
    if not result["missing_required"] and result["has_weights"]:
        result["valid"] = True

    return result


def check_model_integrity(model_dir: Path) -> bool:
    """检查模型文件的完整性"""
    errors = []

    for json_file in model_dir.glob("*.json"):
        try:
            with open(json_file, 'r', encoding='utf-8') as f:
                json.load(f)
        except json.JSONDecodeError as e:
            errors.append(f"{json_file.name}: JSON 解析错误 - {e}")
        except Exception as e:
            errors.append(f"{json_file.name}: 读取错误 - {e}")

    # 显示权重文件大小
    for weight_pattern in ["*.bin", "*.safetensors", "*.onnx", "*.h5"]:
        for weight_file in model_dir.glob(weight_pattern):
            size = weight_file.stat().st_size
            if size == 0:
                errors.append(f"{weight_file.name}: 文件大小为 0")
            elif size > 1024 * 1024 * 1024:
                size_gb = size / (1024 * 1024 * 1024)
                print(f"  {weight_file.name}: {size_gb:.2f} GB")
            elif size > 1024 * 1024:
                size_mb = size / (1024 * 1024)
                print(f"  {weight_file.name}: {size_mb:.2f} MB")
            else:
                size_kb = size / 1024
                print(f"  {weight_file.name}: {size_kb:.2f} KB")

    if errors:
        for err in errors:
            print(f"  ✗ {err}")
        return False
    return True


def test_qwen_inference(model_dir: Path):
    """使用 Qwen 模型进行实际推理测试"""
    print("\n" + "=" * 60)
    print("Qwen 模型推理测试")
    print("=" * 60)

    try:
        from transformers import AutoModelForCausalLM, AutoTokenizer
        import torch
    except ImportError as e:
        print(f"✗ 缺少依赖：{e}")
        print("请运行：pip install transformers torch")
        return False

    print(f"\n模型路径：{model_dir}")
    print("-" * 60)

    # 步骤 1: 加载分词器
    print("\n[1/4] 加载分词器...")
    try:
        tokenizer = AutoTokenizer.from_pretrained(
            str(model_dir),
            trust_remote_code=True
        )
        print(f"    ✓ 分词器加载成功")
        print(f"    - 词表大小：{len(tokenizer):,}")
    except Exception as e:
        print(f"    ✗ 分词器加载失败：{e}")
        return False

    # 步骤 2: 加载模型
    print("\n[2/4] 加载模型...")
    try:
        model = AutoModelForCausalLM.from_pretrained(
            str(model_dir),
            trust_remote_code=True,
            torch_dtype=torch.float16,  # 使用半精度节省显存
            device_map="auto"  # 自动选择设备
        )
        print(f"    ✓ 模型加载成功")
        print(f"    - 参数量：{sum(p.numel() for p in model.parameters()):,}")
        print(f"    - 设备：{model.device}")
    except Exception as e:
        print(f"    ✗ 模型加载失败：{e}")
        return False

    # 步骤 3: 构建测试输入
    print("\n[3/4] 准备测试输入...")
    test_prompts = [
        "你好，请介绍一下你自己。",
        "What is the capital of France?",
    ]

    for prompt in test_prompts:
        print(f"\n    输入：{prompt}")

    # 步骤 4: 执行推理
    print("\n[4/4] 执行推理测试...")
    try:
        for prompt in test_prompts:
            inputs = tokenizer(prompt, return_tensors="pt")
            inputs = inputs.to(model.device)

            with torch.no_grad():
                outputs = model.generate(
                    **inputs,
                    max_new_tokens=50,
                    do_sample=True,
                    temperature=0.7,
                    top_p=0.9,
                )

            response = tokenizer.decode(outputs[0], skip_special_tokens=True)
            # 只显示 prompt + 部分回复
            response_text = response[len(prompt):].strip()
            if len(response_text) > 100:
                response_text = response_text[:100] + "..."
            print(f"    输出：{response_text}")

        print("\n    ✓ 推理测试通过")
        return True

    except Exception as e:
        print(f"    ✗ 推理失败：{e}")
        return False


def main():
    import argparse

    parser = argparse.ArgumentParser(description="模型验证工具")
    parser.add_argument(
        "--model",
        type=str,
        default=None,
        help="指定要验证的模型名称 (默认验证所有)"
    )
    parser.add_argument(
        "--test-inference",
        action="store_true",
        help="对 Qwen 模型进行推理测试"
    )
    args = parser.parse_args()

    print("=" * 60)
    print("模型验证工具")
    print("=" * 60)
    print(f"模型目录：{WEIGHTS_DIR}\n")

    if not WEIGHTS_DIR.exists():
        print("✗ 模型目录不存在")
        return

    model_dirs = [d for d in WEIGHTS_DIR.iterdir() if d.is_dir()]

    if not model_dirs:
        print("未找到任何模型")
        return

    # 过滤指定模型
    if args.model:
        model_dirs = [d for d in model_dirs if args.model in d.name]
        if not model_dirs:
            print(f"✗ 未找到包含 '{args.model}' 的模型")
            return

    print(f"找到 {len(model_dirs)} 个模型:\n")

    all_valid = True

    for model_dir in sorted(model_dirs):
        print("-" * 60)
        print(f"\n模型：{model_dir.name}")
        print(f"路径：{model_dir}")

        result = verify_model(model_dir)

        print(f"\n文件列表 ({len(result['files'])} 个):")
        for f in result['files'][:15]:
            print(f"  - {f}")
        if len(result['files']) > 15:
            print(f"  ... 还有 {len(result['files']) - 15} 个文件")

        print(f"\n验证结果:")
        if result['valid']:
            print(f"  ✓ 模型完整")
        else:
            print(f"  ✗ 模型不完整")
            all_valid = False
            if result['missing_required']:
                print(f"    缺少必需文件：{', '.join(result['missing_required'])}")
            if not result['has_weights']:
                print(f"    未找到权重文件")

        print(f"\n权重文件:")
        if result['weight_files']:
            for wf in result['weight_files']:
                print(f"  - {wf}")
        else:
            print(f"  未找到权重文件")

        print(f"\n完整性检查:")
        if check_model_integrity(model_dir):
            print(f"  ✓ 文件完整性通过")
        else:
            print(f"  ✗ 文件完整性失败")
            all_valid = False

    print("\n" + "=" * 60)
    if all_valid:
        print("✓ 所有模型文件验证通过!")
    else:
        print("✗ 部分模型验证失败")
    print("=" * 60)

    # 推理测试
    if args.test_inference:
        qwen_dir = WEIGHTS_DIR / "Qwen2.5-7B-Instruct"
        if qwen_dir.exists():
            success = test_qwen_inference(qwen_dir)
            if success:
                print("\n" + "=" * 60)
                print("✓ Qwen 模型推理测试成功!")
                print("=" * 60)
            else:
                print("\n" + "=" * 60)
                print("✗ Qwen 模型推理测试失败")
                print("=" * 60)
        else:
            print("\n✗ 未找到 Qwen 模型，跳过推理测试")


if __name__ == "__main__":
    main()