cmp_weights.py

import sys
import gguf

import numpy as np


def find_tensor(name, reader):
    return next((tensor for tensor in reader.tensors if tensor.name == name), None)


def print_tensor(tensor):
    if not isinstance(tensor, list):
        tensor = [tensor]

    for t in tensor:
        print(
            f"{t.n_elements:10} | {list(t.shape)} | {t.tensor_type.name:7} | {t.name}"
        )


def compare_tensors(tensor1, tensor2):
    data1 = np.array(tensor1.data, dtype=np.uint8)
    if isinstance(tensor2, list):
        data2 = np.stack(
            [np.array(t.data, dtype=np.uint8) for t in tensor2],
            axis=0,  # Hmmmmm!
        ).reshape(-1)
    else:
        data2 = np.array(tensor2.data, dtype=np.uint8)
        if tensor2.name.endswith("ffn_gate_inp.weight"):
            data2 = data2.reshape(tensor2.shape)[:, :2].reshape(-1)

    if np.array_equal(data1, data2):
        return True

    indices = np.where(data1 != data2)[0]
    percent_different = 100 * len(indices) / len(data1)

    print(
        tensor1.name,
        f"{percent_different:.2f}%",
        sep="\t",
    )

    return False


def main():
    fn1, fn2 = sys.argv[1:]
    reader1 = gguf.GGUFReader(fn1, "r")
    reader2 = gguf.GGUFReader(fn2, "r")

    print(fn1, fn2)

    success = True
    for tensor1 in reader1.tensors:
        if "_exps" in tensor1.name:
            tensor2 = [
                find_tensor(tensor1.name.replace("_exps", f".{i}"), reader2)
                for i in [0, 1]
            ]
            if any(t is None for t in tensor2):
                tensor2 = None
        else:
            tensor2 = find_tensor(tensor1.name, reader2)
        if tensor2 is None:
            continue
        s = compare_tensors(tensor1, tensor2)
        sys.stdout.write("." if s else "F")
        success &= s

    sys.stdout.write("\n")

    sys.exit(bool(not success))


if __name__ == "__main__":
    main()