| |
| """Rank H3 transformer matrices by mean activation importance.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import re |
| import struct |
| from pathlib import Path |
|
|
| import numpy as np |
|
|
|
|
| FC1 = re.compile(r"model\.diffusion_model\.blocks\.(\d+)\.mlp\.fc1\.weight$") |
|
|
|
|
| def read_entries(path: Path) -> list[dict[str, object]]: |
| rows: list[dict[str, object]] = [] |
| with path.open("rb") as stream: |
| (entry_count,) = struct.unpack("<i", stream.read(4)) |
| for _ in range(entry_count): |
| (name_length,) = struct.unpack("<i", stream.read(4)) |
| name = stream.read(name_length).decode("utf-8") |
| calls, value_count = struct.unpack("<ii", stream.read(8)) |
| values = np.frombuffer(stream.read(value_count * 4), dtype="<f4") |
| match = FC1.fullmatch(name) |
| if match: |
| rows.append( |
| { |
| "block": int(match.group(1)), |
| "calls": calls, |
| "values": value_count, |
| "mean": float(values.mean()), |
| "median": float(np.median(values)), |
| "p95": float(np.quantile(values, 0.95)), |
| "max": float(values.max()), |
| } |
| ) |
| |
| trailer = stream.read() |
| if len(trailer) != 4: |
| raise ValueError(f"unexpected imatrix trailer length: {len(trailer)}") |
| struct.unpack("<i", trailer) |
| return sorted(rows, key=lambda row: float(row["mean"]), reverse=True) |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser() |
| parser.add_argument("imatrix", type=Path) |
| parser.add_argument("--output", type=Path) |
| args = parser.parse_args() |
| rows = read_entries(args.imatrix) |
| rendered = json.dumps(rows, indent=2) + "\n" |
| if args.output: |
| args.output.write_text(rendered, encoding="utf-8") |
| else: |
| print(rendered, end="") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
