| """Упаковка MuScriptor в один самодостаточный архив. |
| |
| Внутри: веса, конфиг, карточка модели, все инструменты и manifest с хешем. |
| Архив намеренно без сжатия: safetensors это уже плотные числа, gzip даёт ~2% |
| выигрыша за десятки минут CPU, а манифест и хеши считаются на лету. |
| |
| Запуск: |
| python tools/pack_bundle.py --model-dir muscriptor-large --out out/muscriptor-bundle.tar |
| """ |
| import argparse |
| import hashlib |
| import io |
| import json |
| import os |
| import struct |
| import sys |
| import tarfile |
| import time |
|
|
| BUNDLE_VERSION = "1.0" |
| CHUNK = 8 * 1024 * 1024 |
|
|
|
|
| def human(size): |
| for unit in ("Б", "КБ", "МБ", "ГБ", "ТБ"): |
| if size < 1024 or unit == "ТБ": |
| return "%.2f %s" % (size, unit) |
| size /= 1024.0 |
|
|
|
|
| def model_stats(path): |
| """Размер, число тензоров и параметров — по заголовку, без чтения данных.""" |
| with open(path, "rb") as handle: |
| length = struct.unpack("<Q", handle.read(8))[0] |
| header = json.loads(handle.read(length)) |
| header.pop("__metadata__", None) |
| params = 0 |
| dtypes = {} |
| for info in header.values(): |
| begin, end = info["data_offsets"] |
| dtypes[info["dtype"]] = dtypes.get(info["dtype"], 0) + 1 |
| params += (end - begin) // 4 |
| return {"tensors": len(header), "parameters": params, "dtypes": dtypes, |
| "bytes": os.path.getsize(path)} |
|
|
|
|
| def add_with_progress(archive, source, arcname, digest, total, label): |
| """Кладёт файл в tar потоком, попутно считая sha256 и печатая проценты.""" |
| size = os.path.getsize(source) |
| info = tarfile.TarInfo(arcname) |
| info.size = size |
| info.mtime = int(time.time()) |
| with open(source, "rb") as handle: |
| reader = _HashingReader(handle, digest, total, label) |
| archive.addfile(info, reader) |
|
|
|
|
| class _HashingReader(io.RawIOBase): |
| def __init__(self, handle, digest, total, label): |
| self.handle, self.digest, self.total, self.label = handle, digest, total, label |
| self.done = 0 |
| self.last = -1 |
|
|
| def readable(self): |
| return True |
|
|
| def read(self, size=-1): |
| block = self.handle.read(size if size and size > 0 else CHUNK) |
| if block: |
| self.digest.update(block) |
| self.done += len(block) |
| if self.total: |
| percent = int(self.done * 100 / self.total) |
| if percent != self.last: |
| self.last = percent |
| print("\r %s: %3d%% %s / %s" |
| % (self.label, percent, human(self.done), human(self.total)), |
| end="", flush=True) |
| elif self.total: |
| print("\r %s: 100%% %s" % (self.label, human(self.total)), flush=True) |
| return block |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--model-dir", default="muscriptor-large") |
| parser.add_argument("--out", default="out/muscriptor-large-bundle.tar") |
| parser.add_argument("--tools", default=None, help="папка с инструментами (по умолчанию рядом с этим скриптом)") |
| parser.add_argument("--license-note", default=None) |
| args = parser.parse_args() |
|
|
| weights = os.path.join(args.model_dir, "model.safetensors") |
| config = os.path.join(args.model_dir, "config.json") |
| if not os.path.isfile(weights): |
| raise SystemExit("нет весов: " + weights) |
|
|
| print("[5%] считаю параметры по заголовку") |
| stats = model_stats(weights) |
| print(" тензоров %d, параметров %.1f млн, типы %s, вес %s" |
| % (stats["tensors"], stats["parameters"] / 1e6, stats["dtypes"], human(stats["bytes"]))) |
|
|
| tools_dir = args.tools or os.path.dirname(os.path.abspath(__file__)) |
| tools = sorted(name for name in os.listdir(tools_dir) if name.endswith(".py")) |
| print("[10%] инструменты: " + ", ".join(tools) if tools else "[10%] инструменты не найдены") |
|
|
| os.makedirs(os.path.dirname(os.path.abspath(args.out)), exist_ok=True) |
| digest = hashlib.sha256() |
| manifest = { |
| "bundle_version": BUNDLE_VERSION, |
| "built_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), |
| "model": { |
| "name": "MuScriptor Large", |
| "source": "https://huggingface.co/MuScriptor/muscriptor-large", |
| "variant": "large", |
| "tensors": stats["tensors"], |
| "parameters": stats["parameters"], |
| "dtypes": stats["dtypes"], |
| "weights_bytes": stats["bytes"], |
| }, |
| "license": { |
| "weights": "CC-BY-NC-4.0", |
| "code": "MIT", |
| "note": "Веса некоммерческие: их нельзя использовать в платном продукте.", |
| }, |
| "files": ["model/model.safetensors", "model/config.json", "model/README.md", "manifest.json"], |
| "tools": tools, |
| } |
|
|
| print("[15%] пишу архив " + args.out) |
| with tarfile.open(args.out, "w") as archive: |
| for name in tools: |
| path = os.path.join(tools_dir, name) |
| info = tarfile.TarInfo("tools/" + name) |
| info.size = os.path.getsize(path) |
| info.mtime = int(time.time()) |
| with open(path, "rb") as handle: |
| archive.addfile(info, handle) |
| manifest["files"].append("tools/" + name) |
| if args.license_note and os.path.isfile(args.license_note): |
| data = open(args.license_note, "rb").read() |
| info = tarfile.TarInfo("LICENSE_NOTICE.txt") |
| info.size = len(data) |
| info.mtime = int(time.time()) |
| archive.addfile(info, io.BytesIO(data)) |
| manifest["files"].append("LICENSE_NOTICE.txt") |
|
|
| if os.path.isfile(config): |
| data = open(config, "rb").read() |
| info = tarfile.TarInfo("model/config.json") |
| info.size = len(data) |
| info.mtime = int(time.time()) |
| archive.addfile(info, io.BytesIO(data)) |
| card = os.path.join(args.model_dir, "README.md") |
| if os.path.isfile(card): |
| data = open(card, "rb").read() |
| info = tarfile.TarInfo("model/README.md") |
| info.size = len(data) |
| info.mtime = int(time.time()) |
| archive.addfile(info, io.BytesIO(data)) |
|
|
| add_with_progress(archive, weights, "model/model.safetensors", digest, |
| stats["bytes"], "веса") |
|
|
| manifest["model"]["sha256"] = digest.hexdigest() |
| blob = json.dumps(manifest, ensure_ascii=False, indent=2).encode("utf-8") |
| info = tarfile.TarInfo("manifest.json") |
| info.size = len(blob) |
| info.mtime = int(time.time()) |
| archive.addfile(info, io.BytesIO(blob)) |
|
|
| size = os.path.getsize(args.out) |
| print("[100%%] готово: %s (%s)" % (args.out, human(size))) |
| print(" sha256 весов: " + manifest["model"]["sha256"]) |
| print(" в истории: python tools/pack_bundle.py --model-dir ... --out ...") |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| sys.exit(main()) |
|
|
