File size: 2,106 Bytes
d18b291 | 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 | #!/usr/bin/env python3
"""Emit a compact, reproducible tensor and parameter audit for GGUF files."""
from __future__ import annotations
import argparse
import hashlib
import json
from collections import Counter
from pathlib import Path
from gguf import GGUFReader
def sha256(path: Path) -> str:
digest = hashlib.sha256()
with path.open("rb") as stream:
for chunk in iter(lambda: stream.read(8 * 1024 * 1024), b""):
digest.update(chunk)
return digest.hexdigest()
def audit(path: Path, with_hash: bool) -> dict[str, object]:
reader = GGUFReader(path, "r")
tensors: Counter[str] = Counter()
parameters: Counter[str] = Counter()
protected: dict[str, str] = {}
for tensor in reader.tensors:
tensor_type = tensor.tensor_type.name.lower()
tensors[tensor_type] += 1
parameters[tensor_type] += int(tensor.n_elements)
if tensor.name.endswith(
(
"audio_patch_proj.weight",
"video_patch_proj.weight",
"condition_proj.weight",
)
):
protected[tensor.name] = tensor_type
result: dict[str, object] = {
"file": path.name,
"bytes": path.stat().st_size,
"tensor_count": len(reader.tensors),
"tensor_histogram": dict(sorted(tensors.items())),
"parameter_histogram": dict(sorted(parameters.items())),
"protected_tensors": dict(sorted(protected.items())),
}
if with_hash:
result["sha256"] = sha256(path)
return result
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("files", nargs="+", type=Path)
parser.add_argument("--sha256", action="store_true")
parser.add_argument("--output", type=Path)
args = parser.parse_args()
report = [audit(path.resolve(), args.sha256) for path in args.files]
rendered = json.dumps(report, indent=2, sort_keys=True) + "\n"
if args.output:
args.output.write_text(rendered, encoding="utf-8")
else:
print(rendered, end="")
if __name__ == "__main__":
main()
|