#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""patch_imx_adapter.py — 远端 r13_chain/r13_refine 的 imatrix 新格式适配补丁
新 llama.cpp imatrix = <name>.weight.in_sum2 (1-D f32 列激活能量) + .counts
S3/S5 原代码只认 (in,in) 全矩阵 → 全部落 identity, imatrix 加权失效。
补丁: 查找链加 .in_sum2 候选 + 1-D→torch.diag 对角矩阵 (数学上等价 GPTQ 对角 Hessian)。
锚点断言 count==1 + compile 语法验证, 双保险。
"""
A_OLD = '''        hk = base if base in imx else (base + ".weight" if base + ".weight" in imx else None)
        if hk is not None:
            hn = np.ascontiguousarray(imx[hk].data).astype(np.float32)
            H = None
            if hn.shape == (shape[1], shape[1]):
                H = torch.from_numpy(hn).to(dev)'''
A_NEW = '''        hk = None
        for cd in (base, base + ".weight", base + ".weight.in_sum2", base + ".in_sum2"):
            if cd in imx:
                hk = cd
                break
        if hk is not None:
            hn = np.ascontiguousarray(imx[hk].data).astype(np.float32)
            H = None
            if hn.ndim == 1 and hn.shape[0] == shape[1]:
                H = torch.diag(torch.from_numpy(hn).to(dev) + 1e-8)
            elif hn.shape == (shape[1], shape[1]):
                H = torch.from_numpy(hn).to(dev)'''
B_OLD = '''    hn = np.ascontiguousarray(imx_t.data).astype(np.float32)
    if hn.shape != (in_dim, in_dim):
        return None'''
B_NEW = '''    hn = np.ascontiguousarray(imx_t.data).astype(np.float32)
    if hn.ndim == 1 and hn.shape[0] == in_dim:
        hn = np.diag(hn + 1e-8)
    if hn.shape != (in_dim, in_dim):
        return None'''
C_OLD = '''        hk = base if base in imx else (base + ".weight" if base + ".weight" in imx else None)'''
C_NEW = '''        hk = None
        for cd in (base, base + ".weight", base + ".weight.in_sum2", base + ".in_sum2"):
            if cd in imx:
                hk = cd
                break'''

def main():
    for path, olds, news in (
        ("/tmp/k8b/r13_chain.py", [A_OLD], [A_NEW]),
        ("/tmp/k8b/r13_refine.py", [B_OLD, C_OLD], [B_NEW, C_NEW]),
    ):
        src = open(path).read()
        for o, n in zip(olds, news):
            cnt = src.count(o)
            assert cnt == 1, "anchor not unique in %s: count=%d" % (path, cnt)
            src = src.replace(o, n, 1)
        compile(src, path, "exec")
        open(path, "w").write(src)
        print("PATCHED", path)
    print("PATCH_OK imx in_sum2 diag adapter")
main()