#!/usr/bin/env python3
"""D-57 · llama.cpp's imatrix.gguf → the importance safetensors the trainer's N-bit scale search reads
(`quant.importance`).

    python scripts/imatrix_to_importance.py /path/to/data/ptq/imatrix.gguf \
        /path/to/data/ptq/importance.safetensors
"""
from __future__ import annotations

import sys
from pathlib import Path

REPO = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(REPO / "src"))
from lowbit.export.gguf import gguf_module  # noqa: E402
from lowbit.quant.importance import from_imatrix  # noqa: E402

LLAMA = Path("/path/to/llama.cpp")


def main() -> None:
    from safetensors.torch import save_file
    src, dst = sys.argv[1], sys.argv[2]
    imp = from_imatrix(src, gguf_module(LLAMA))
    save_file(imp, dst, metadata={"source": str(src)})
    layers = len(imp) // 2
    print(f"{dst}: {layers} layers, gate_up {tuple(imp['layers.0.gate_up'].shape)}, "
          f"down {tuple(imp['layers.0.down'].shape)}")


if __name__ == "__main__":
    main()