| """Проверка весов MuScriptor: заголовок safetensors, конфиг, число параметров, sha256. |
| |
| Запуск: |
| python tools/verify_model.py --model-dir muscriptor-large [--hash] |
| |
| Ничего не загружает в память целиком: читает только 8-байтовую длину заголовка |
| и сам JSON-заголовок, поэтому работает и на 5 ГБ, и на 50 ГБ. |
| """ |
| import argparse |
| import hashlib |
| import json |
| import os |
| import struct |
| import sys |
|
|
| DTYPE_BYTES = {"F64": 8, "F32": 4, "F16": 2, "BF16": 2, "I64": 8, "I32": 4, |
| "I16": 2, "I8": 1, "U8": 1, "BOOL": 1} |
|
|
|
|
| def human(size): |
| for unit in ("Б", "КБ", "МБ", "ГБ", "ТБ"): |
| if size < 1024 or unit == "ТБ": |
| return "%.2f %s" % (size, unit) |
| size /= 1024.0 |
|
|
|
|
| def read_header(path): |
| with open(path, "rb") as handle: |
| raw = handle.read(8) |
| if len(raw) != 8: |
| raise SystemExit("файл короче 8 байт: " + path) |
| length = struct.unpack("<Q", raw)[0] |
| if length <= 0 or length > 200 * 1024 * 1024: |
| raise SystemExit("подозрительная длина заголовка: %d" % length) |
| header = json.loads(handle.read(length)) |
| meta = header.pop("__metadata__", None) |
| return header, meta, 8 + length |
|
|
|
|
| def sha256(path, chunk=8 * 1024 * 1024, label="sha256"): |
| digest = hashlib.sha256() |
| total = os.path.getsize(path) |
| done = 0 |
| with open(path, "rb") as handle: |
| while True: |
| block = handle.read(chunk) |
| if not block: |
| break |
| digest.update(block) |
| done += len(block) |
| if total: |
| print("\r %s: %5.1f%%" % (label, done * 100.0 / total), end="", flush=True) |
| print("\r %s: 100.0%% %s" % (label, digest.hexdigest())) |
| return digest.hexdigest() |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--model-dir", default="muscriptor-large") |
| parser.add_argument("--hash", action="store_true", help="посчитать sha256 (долго)") |
| args = parser.parse_args() |
|
|
| weights = os.path.join(args.model_dir, "model.safetensors") |
| config_path = os.path.join(args.model_dir, "config.json") |
| if not os.path.isfile(weights): |
| raise SystemExit("нет файла весов: " + weights) |
|
|
| print("[10%] читаю заголовок safetensors") |
| tensors, meta, data_start = read_header(weights) |
| file_size = os.path.getsize(weights) |
|
|
| dtypes = {} |
| total_bytes = 0 |
| for info in tensors.values(): |
| dtypes[info["dtype"]] = dtypes.get(info["dtype"], 0) + 1 |
| begin, end = info["data_offsets"] |
| total_bytes += end - begin |
|
|
| expected = data_start + total_bytes |
| print("[30%%] файл: %s | данные: %s | заголовок: %s" |
| % (human(file_size), human(total_bytes), human(data_start))) |
| print(" тензоров: %d | типы: %s" % (len(tensors), dtypes)) |
| print(" параметров: %.1f млн" % (total_bytes / 4.0 / 1e6)) |
|
|
| if expected != file_size: |
| print("[!] РАСХОЖДЕНИЕ: ожидалось %s, в файле %s" % (human(expected), human(file_size))) |
| ok_size = False |
| else: |
| ok_size = True |
| print("[50%] размеры сходятся байт в байт") |
|
|
| config = {} |
| if os.path.isfile(config_path): |
| with open(config_path, encoding="utf-8") as handle: |
| config = json.load(handle) |
| print("[60%] config.json: " + json.dumps(config, ensure_ascii=False)) |
| dim = int(config.get("dim") or 0) |
| layers = int(config.get("num_layers") or 0) |
| if dim: |
| probe = "linears.0.weight" |
| if probe in tensors: |
| out_features = tensors[probe]["shape"][0] |
| print(" выходная голова: %s -> %d токенов" % (probe, out_features)) |
| hits = sum(1 for name in tensors if name.startswith("layers.")) |
| print(" тензоров в layers.*: %d (конфиг заявляет %d слоёв)" % (hits, layers)) |
| else: |
| print("[!] нет config.json рядом с весами") |
|
|
| if meta: |
| print("[70%] метаданные: " + json.dumps(meta, ensure_ascii=False)[:400]) |
|
|
| exit_code = 0 if ok_size else 1 |
| if args.hash: |
| print("[80%] считаю sha256 (это долго, ~1-2 минуты на 5 ГБ)") |
| sha256(weights) |
| else: |
| print("[80%] sha256 пропущен (добавь --hash)") |
|
|
| print("[100%] итог: " + ("ГОДНО" if exit_code == 0 else "ФАЙЛ ПОВРЕЖДЁН")) |
| return exit_code |
|
|
|
|
| if __name__ == "__main__": |
| sys.exit(main()) |
|
|
