File size: 4,809 Bytes
ca6d1b8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 | """Проверка весов 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())
|