kg-agent/scripts/download_model.py

"""Download and validate the configured local embedding model; never sends business data."""
from __future__ import annotations
import argparse
import json
import os
from pathlib import Path
import shutil
import subprocess
import sys

from step2_graph_building.config import get_embedding_model_dir

ROOT = Path(__file__).resolve().parents[1]


def model_files_ready(path: Path) -> bool:
    required = ['config.json', 'modules.json', 'tokenizer.json', 'tokenizer_config.json']
    try:
        for name in required:
            if not (path / name).is_file() or (path / name).stat().st_size == 0:
                return False
            json.loads((path / name).read_text(encoding='utf-8'))
        modules = json.loads((path / 'modules.json').read_text(encoding='utf-8'))
        for module in modules:
            relative = module.get('path', '')
            if (relative and module.get('type') != 'sentence_transformers.models.Normalize'
                    and not (path / relative / 'config.json').is_file()):
                return False
        index = path / 'model.safetensors.index.json'
        if index.is_file():
            weights = set(json.loads(index.read_text(encoding='utf-8'))['weight_map'].values())
        else:
            weights = {'model.safetensors'}
        return bool(weights) and all((path / name).is_file() and (path / name).stat().st_size > 0 for name in weights)
    except (OSError, ValueError, KeyError, TypeError, AttributeError):
        return False


def verify_model(path: Path) -> None:
    import numpy as np
    from sentence_transformers import SentenceTransformer
    model = SentenceTransformer(str(path), device='cpu', local_files_only=True)
    vectors = model.encode(['部署检查'], normalize_embeddings=True, show_progress_bar=False)
    if len(vectors) != 1 or not np.isfinite(vectors).all() or np.linalg.norm(vectors) == 0:
        raise RuntimeError('本地模型编码验证失败')


def ensure_model(*, check_only=False, force=False):
    path = get_embedding_model_dir()
    repo = os.getenv('EMBEDDING_MODEL_REPO', 'Qwen/Qwen3-Embedding-0.6B').strip()
    source = os.getenv('EMBEDDING_MODEL_SOURCE', 'modelscope').strip().lower()
    if not repo or source not in {'modelscope', 'huggingface'}:
        raise ValueError('请配置有效 EMBEDDING_MODEL_REPO，SOURCE 仅支持 modelscope/huggingface')
    if force or not model_files_ready(path):
        if check_only:
            raise RuntimeError(f'模型文件不完整: {path}')
        print(f'下载模型 {repo} → {path}（{source}）', flush=True)
        path.mkdir(parents=True, exist_ok=True)
        if source == 'modelscope':
            uv = shutil.which('uv')
            if not uv:
                raise RuntimeError('ModelScope 下载需要 uv；请先运行 sh deploy.sh setup')
            # ModelScope lives in an ephemeral uv environment, not project dependencies.
            code = 'from modelscope import snapshot_download; import sys; snapshot_download(sys.argv[1], local_dir=sys.argv[2])'
            subprocess.run([uv, 'run', '--frozen', '--with', 'modelscope', 'python', '-c', code, repo, str(path)], cwd=ROOT, check=True)
        else:
            from huggingface_hub import snapshot_download
            snapshot_download(repo_id=repo, local_dir=str(path), force_download=force)
        if not model_files_ready(path):
            raise RuntimeError(f'模型下载未完整，请检查网络后重试: {path}')
    else:
        print(f'模型文件已存在，跳过下载: {path}', flush=True)
    verify_model(path)
    print('本地模型加载和编码检查通过。', flush=True)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--check-only', action='store_true', help='仅校验本地模型，禁止下载')
    parser.add_argument('--force', action='store_true', help='重新调用下载器修复模型文件')
    args = parser.parse_args()
    if args.check_only and args.force:
        parser.error('--check-only 与 --force 不能同时使用')
    ensure_model(check_only=args.check_only, force=args.force)


if __name__ == '__main__':
    main()