144 lines
6.6 KiB
Python
144 lines
6.6 KiB
Python
"""Prepare immutable model/tokenizer assets during build, never during startup."""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
import os
|
|
import re
|
|
import shutil
|
|
import subprocess
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
DEFAULT_MODEL = 'intfloat/multilingual-e5-small'
|
|
DEFAULT_REVISION = '614241f622f53c4eeff9890bdc4f31cfecc418b3'
|
|
DEFAULT_ROOT = Path(__file__).resolve().parents[2] / 'runtime_assets'
|
|
|
|
|
|
def asset_root() -> Path:
|
|
return Path(os.getenv('ZP_RUNTIME_ASSETS', str(DEFAULT_ROOT)))
|
|
|
|
|
|
def manifest(root: Path | None = None) -> dict | None:
|
|
path = (root or asset_root()) / 'manifest.json'
|
|
if not path.exists():
|
|
if os.getenv('ZP_ASSETS_REQUIRED') == '1':
|
|
raise RuntimeError('Chýbajú runtime assets; zostavte image alebo spustite python -m tools.runtime.runtime_assets prepare.')
|
|
return None
|
|
return json.loads(path.read_text())
|
|
|
|
|
|
def configure_tokenizer() -> None:
|
|
data = manifest()
|
|
if data is not None:
|
|
# Bundled immutable files take precedence over an empty user cache.
|
|
os.environ['TIKTOKEN_CACHE_DIR'] = str(asset_root() / 'tiktoken')
|
|
|
|
|
|
def local_model(model_name: str, revision: str | None) -> str | None:
|
|
data = manifest()
|
|
if data and data['embedding']['model'] == model_name:
|
|
if revision and revision != data['embedding']['revision']:
|
|
raise RuntimeError('Požadovaná revízia modelu nie je v image; vytvorte nový image s touto revíziou.')
|
|
return str(asset_root() / 'model')
|
|
if data and os.getenv('ZP_ASSETS_REQUIRED') == '1':
|
|
raise RuntimeError('Požadovaný embeddingový model nie je v image; vytvorte nový image.')
|
|
return None
|
|
|
|
|
|
def verify(root: Path) -> dict:
|
|
data = json.loads((root / 'manifest.json').read_text())
|
|
for relative, digest in data['files'].items():
|
|
path = root / relative
|
|
if not path.resolve().is_relative_to(root.resolve()):
|
|
raise RuntimeError('Manifest obsahuje cestu mimo runtime assets')
|
|
if not path.is_file() or hashlib.sha256(path.read_bytes()).hexdigest() != digest:
|
|
raise RuntimeError(f'Chýbajúci alebo poškodený runtime asset: {relative}')
|
|
if not any(name.startswith('tiktoken/') for name in data['files']):
|
|
raise RuntimeError('Manifest neobsahuje tokenizér')
|
|
required = {'model/config.json', 'model/modules.json', 'model/tokenizer.json'}
|
|
if not required.issubset(data['files']) or not any(name.endswith('.safetensors') for name in data['files']):
|
|
raise RuntimeError('Manifest neobsahuje kompletný embeddingový model')
|
|
return data
|
|
|
|
|
|
def offline_probe(root: Path) -> None:
|
|
"""Fresh interpreter, with network calls forbidden even on an online builder."""
|
|
environment = {**os.environ, 'ZP_RUNTIME_ASSETS': str(root.resolve()), 'ZP_ASSETS_REQUIRED': '1',
|
|
'HF_HUB_OFFLINE': '1', 'TRANSFORMERS_OFFLINE': '1', 'HF_HOME': str(root / 'unused-user-cache'),
|
|
'TIKTOKEN_CACHE_DIR': str(root / 'unused-token-cache')}
|
|
code = '''
|
|
import socket
|
|
def denied(*args, **kwargs): raise RuntimeError('Network access during offline asset probe')
|
|
socket.socket.connect = denied
|
|
socket.socket.connect_ex = denied
|
|
socket.getaddrinfo = denied
|
|
from tools.runtime.runtime_assets import configure_tokenizer, manifest, local_model
|
|
configure_tokenizer()
|
|
import tiktoken
|
|
for encoding in manifest()['encodings']:
|
|
assert tiktoken.get_encoding(encoding).encode('Overenie štartu bez siete.')
|
|
from sentence_transformers import SentenceTransformer
|
|
data = manifest()['embedding']
|
|
model = SentenceTransformer(local_model(data['model'], data['revision']), local_files_only=True)
|
|
assert model.encode(['query: overenie bez siete'], show_progress_bar=False).shape[0] == 1
|
|
print('Offline model/tokenizer probe: OK')
|
|
'''
|
|
subprocess.run([sys.executable, '-c', code], env=environment, check=True)
|
|
|
|
|
|
def prepare(root: Path, model: str, revision: str | None, encodings: list[str]) -> dict:
|
|
revision = revision or (DEFAULT_REVISION if model == DEFAULT_MODEL else '')
|
|
if not re.fullmatch('[0-9a-f]{40}', revision):
|
|
raise ValueError('Príprava modelu vyžaduje presnú 40-znakovú Git revíziu')
|
|
if (root / 'manifest.json').is_file():
|
|
data = verify(root)
|
|
if data['embedding'] == {'model': model, 'revision': revision} and data['encodings'] == encodings:
|
|
offline_probe(root)
|
|
return data
|
|
raise RuntimeError('Assets patria inej konfigurácii; použite nový/prázdny pracovný adresár.')
|
|
from huggingface_hub import snapshot_download
|
|
from huggingface_hub.errors import LocalEntryNotFoundError
|
|
try:
|
|
source = snapshot_download(model, revision=revision, local_files_only=True)
|
|
except (FileNotFoundError, LocalEntryNotFoundError):
|
|
source = snapshot_download(model, revision=revision,
|
|
allow_patterns=['*.json', '*.safetensors', '*.model', '*.txt'],
|
|
ignore_patterns=['onnx/*', 'openvino/*'])
|
|
root.mkdir(parents=True, exist_ok=True)
|
|
shutil.copytree(source, root / 'model', dirs_exist_ok=True, symlinks=False,
|
|
ignore=shutil.ignore_patterns('onnx', 'openvino', '*.bin', '*.md', '.cache', '.gitattributes'))
|
|
environment = {**os.environ, 'TIKTOKEN_CACHE_DIR': str(root.resolve() / 'tiktoken')}
|
|
subprocess.run([sys.executable, '-c',
|
|
'import sys,tiktoken; [tiktoken.get_encoding(name) for name in sys.argv[1:]]', *encodings],
|
|
env=environment, check=True)
|
|
files = {str(p.relative_to(root)): hashlib.sha256(p.read_bytes()).hexdigest()
|
|
for folder in ['model', 'tiktoken'] for p in sorted((root / folder).rglob('*')) if p.is_file()}
|
|
data = {'schema_version': 1, 'embedding': {'model': model, 'revision': revision},
|
|
'encodings': encodings, 'files': files}
|
|
(root / 'manifest.json').write_text(json.dumps(data, sort_keys=True, indent=2) + '\n')
|
|
verify(root)
|
|
offline_probe(root)
|
|
return data
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument('command', choices=['prepare', 'verify'])
|
|
parser.add_argument('--output', type=Path, default=asset_root())
|
|
parser.add_argument('--model', default=DEFAULT_MODEL)
|
|
parser.add_argument('--revision')
|
|
parser.add_argument('--encoding', action='append')
|
|
args = parser.parse_args()
|
|
if args.command == 'prepare':
|
|
data = prepare(args.output, args.model, args.revision, list(dict.fromkeys(args.encoding or ['cl100k_base'])))
|
|
else:
|
|
data = verify(args.output)
|
|
offline_probe(args.output)
|
|
print(json.dumps({'model': data['embedding'], 'encodings': data['encodings'], 'files': len(data['files'])}))
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|