290 lines
11 KiB
Python
290 lines
11 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
|
|
|
|
|
|
RERANKER_MODEL = "cross-encoder/mmarco-mMiniLMv2-L12-H384-v1"
|
|
RERANKER_REVISION = "1427fd652930e4ba29e8149678df786c240d8825"
|
|
STANZA_LANGUAGE = "sk"
|
|
STANZA_PROCESSORS = "tokenize,pos,lemma"
|
|
|
|
|
|
def _hashes(root: Path, folders) -> dict:
|
|
return {
|
|
str(p.relative_to(root)): hashlib.sha256(p.read_bytes()).hexdigest()
|
|
for folder in folders
|
|
for p in sorted((root / folder).rglob("*"))
|
|
if p.is_file()
|
|
}
|
|
|
|
|
|
def verify_extras(root: Path) -> dict:
|
|
"""Reranker and Stanza Slovak files unchanged since preparation."""
|
|
path = root / "extras.json"
|
|
if not path.is_file():
|
|
raise RuntimeError("Chýba manifest doplnkových modelov (reranker, Stanza).")
|
|
data = json.loads(path.read_text())
|
|
if _hashes(root, ["reranker", "stanza"]) != data["files"]:
|
|
raise RuntimeError("Doplnkové modely sa líšia od manifestu.")
|
|
return data
|
|
|
|
|
|
def extras_probe(root: Path) -> None:
|
|
"""Load both models without network access."""
|
|
code = f"""
|
|
import socket
|
|
def denied(*args, **kwargs): raise RuntimeError('Network access during offline asset probe')
|
|
socket.create_connection = denied
|
|
from sentence_transformers import CrossEncoder
|
|
model = CrossEncoder({str(root / 'reranker')!r}, local_files_only=True, device='cpu')
|
|
assert len(model.predict([('otázka', 'text')])) == 1
|
|
import stanza
|
|
nlp = stanza.Pipeline({STANZA_LANGUAGE!r}, dir={str(root / 'stanza')!r}, processors={STANZA_PROCESSORS!r},
|
|
download_method=None, verbose=False, use_gpu=False)
|
|
assert nlp('chatbotom').sentences[0].words[0].lemma == 'chatbot'
|
|
print('Offline reranker/Stanza probe: OK')
|
|
"""
|
|
environment = {**os.environ, "HF_HUB_OFFLINE": "1", "TRANSFORMERS_OFFLINE": "1"}
|
|
subprocess.run([sys.executable, "-c", code], env=environment, check=True)
|
|
|
|
|
|
def prepare_extras(root: Path) -> dict:
|
|
"""Pinned cross-encoder and Stanza Slovak lemmatizer, copied into root."""
|
|
if (root / "extras.json").is_file():
|
|
data = verify_extras(root)
|
|
if data["reranker"] == {"model": RERANKER_MODEL, "revision": RERANKER_REVISION}:
|
|
extras_probe(root)
|
|
return data
|
|
raise RuntimeError("Doplnkové modely patria inej konfigurácii.")
|
|
from huggingface_hub import snapshot_download
|
|
from huggingface_hub.errors import LocalEntryNotFoundError
|
|
try:
|
|
source = snapshot_download(RERANKER_MODEL, revision=RERANKER_REVISION, local_files_only=True)
|
|
except (FileNotFoundError, LocalEntryNotFoundError):
|
|
source = snapshot_download(RERANKER_MODEL, revision=RERANKER_REVISION,
|
|
allow_patterns=["*.json", "*.safetensors", "*.model", "*.txt"])
|
|
root.mkdir(parents=True, exist_ok=True)
|
|
shutil.copytree(source, root / "reranker", dirs_exist_ok=True, symlinks=False,
|
|
ignore=shutil.ignore_patterns("onnx", "openvino", "*.bin", "*.md", ".cache", ".gitattributes"))
|
|
import stanza
|
|
from stanza.resources.common import DEFAULT_MODEL_DIR
|
|
cached = Path(DEFAULT_MODEL_DIR)
|
|
if (cached / STANZA_LANGUAGE).is_dir() and (cached / "resources.json").is_file():
|
|
(root / "stanza").mkdir(parents=True, exist_ok=True)
|
|
shutil.copy2(cached / "resources.json", root / "stanza" / "resources.json")
|
|
shutil.copytree(cached / STANZA_LANGUAGE, root / "stanza" / STANZA_LANGUAGE, dirs_exist_ok=True)
|
|
else:
|
|
stanza.download(STANZA_LANGUAGE, model_dir=str(root / "stanza"), processors=STANZA_PROCESSORS, verbose=False)
|
|
data = {
|
|
"schema_version": 1,
|
|
"reranker": {"model": RERANKER_MODEL, "revision": RERANKER_REVISION},
|
|
"stanza": {"version": stanza.__version__, "language": STANZA_LANGUAGE, "processors": STANZA_PROCESSORS},
|
|
"files": _hashes(root, ["reranker", "stanza"]),
|
|
}
|
|
(root / "extras.json").write_text(json.dumps(data, sort_keys=True, indent=2) + "\n")
|
|
extras_probe(root)
|
|
return data
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("command", choices=["prepare", "verify", "prepare-extras", "verify-extras"])
|
|
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 in {"prepare-extras", "verify-extras"}:
|
|
data = prepare_extras(args.output) if args.command == "prepare-extras" else verify_extras(args.output)
|
|
if args.command == "verify-extras":
|
|
extras_probe(args.output)
|
|
print(json.dumps({"reranker": data["reranker"], "stanza": data["stanza"], "files": len(data["files"])}))
|
|
return
|
|
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()
|