business-assistant/scripts/deploy_models.py

"""Download deployment assets, never model code; run with the locked project venv."""
import argparse
import json
from pathlib import Path

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


def bge_valid(path):
    try:
        config = json.loads((path / 'config.json').read_text(encoding='utf-8'))
        if config.get('model_type') != 'bert' or config.get('hidden_size') != 1024:
            return False
        from transformers import AutoModel, AutoTokenizer
        AutoTokenizer.from_pretrained(path, local_files_only=True, trust_remote_code=False)
        model = AutoModel.from_pretrained(path, local_files_only=True, trust_remote_code=False)
        del model
        return True
    except Exception:
        return False


def tokenizer_valid(path):
    try:
        from tokenizers import Tokenizer
        Tokenizer.from_file(str(path))
        return True
    except Exception:
        return False


def download(check_only=False):
    from huggingface_hub import snapshot_download, hf_hub_download
    model = ROOT / 'models/bge-large-zh-v1.5'
    tokenizer = ROOT / 'models/deepseek-tokenizer/tokenizer.json'
    model_ready = bge_valid(model)
    if check_only:
        if not model_ready or not tokenizer_valid(tokenizer):
            raise ValueError('active_service_model_assets_invalid_stop_before_repair')
        print('Existing model assets validated without modification.', flush=True)
        return
    if not model_ready:
        snapshot_download('BAAI/bge-large-zh-v1.5', local_dir=model,
                          allow_patterns=['config.json', 'tokenizer.json', 'tokenizer_config.json',
                                          'special_tokens_map.json', 'vocab.txt', 'pytorch_model.bin'])
        model_ready = bge_valid(model)
    if not model_ready:
        raise ValueError('bge_model_validation_failed')
    if not tokenizer_valid(tokenizer):
        hf_hub_download('deepseek-ai/DeepSeek-V3', 'tokenizer.json', local_dir=tokenizer.parent)
    if not tokenizer_valid(tokenizer):
        raise ValueError('summary_tokenizer_validation_failed')
    print('BGE model and summary tokenizer ready.', flush=True)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--check-only', action='store_true')
    args = parser.parse_args()
    try:
        download(args.check_only)
    except Exception as exc:
        print('model_download_failed:' + type(exc).__name__)
        raise SystemExit(1) from None