RWKV-Runner/backend-python/convert_safetensors.py at master · josStorer/RWKV-Runner · GitHub

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
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
import collections
import numpy
import os
import torch
from safetensors.torch import serialize_file, load_file

import argparse

parser = argparse.ArgumentParser()
parser.add_argument("--input", type=str, help="Path to input pth model")
parser.add_argument(
    "--output",
    type=str,
    default="./converted.st",
    help="Path to output safetensors model",
)
args = parser.parse_args()


def rename_key(rename, name):
    for k, v in rename.items():
        if k in name:
            name = name.replace(k, v)
    return name


def convert_file(pt_filename: str, sf_filename: str, rename={}, transpose_names=[]):
    loaded: collections.OrderedDict = torch.load(pt_filename, map_location="cpu")
    if "state_dict" in loaded:
        loaded = loaded["state_dict"]

    kk = list(loaded.keys())
    version = 4
    for x in kk:
        if "ln_x" in x:
            version = max(5, version)
        if "gate.weight" in x:
            version = max(5.1, version)
        if int(version) == 5 and "att.time_decay" in x:
            if len(loaded[x].shape) > 1:
                if loaded[x].shape[1] > 1:
                    version = max(5.2, version)
        if "time_maa" in x:
            version = max(6, version)

    print(f"Model detected: v{version:.1f}")

    if version == 5.1:
        _, n_emb = loaded["emb.weight"].shape
        for k in kk:
            if "time_decay" in k or "time_faaaa" in k:
                # print(k, mm[k].shape)
                loaded[k] = (
                    loaded[k].unsqueeze(1).repeat(1, n_emb // loaded[k].shape[0])
                )

    with torch.no_grad():
        for k in kk:
            new_k = rename_key(rename, k).lower()
            v = loaded[k].half()
            del loaded[k]
            for transpose_name in transpose_names:
                if transpose_name in new_k:
                    dims = len(v.shape)
                    v = v.transpose(dims - 2, dims - 1)
                    break
            print(f"{new_k}\t{v.shape}\t{v.dtype}")
            loaded[new_k] = {
                "dtype": str(v.dtype).split(".")[-1],
                "shape": v.shape,
                "data": v.numpy().tobytes(),
            }

    dirname = os.path.dirname(sf_filename)
    os.makedirs(dirname, exist_ok=True)
    serialize_file(loaded, sf_filename, metadata={"format": "pt"})
    # reloaded = load_file(sf_filename)
    # for k in loaded:
    #     pt_tensor = torch.Tensor(
    #         numpy.frombuffer(
    #             bytearray(loaded[k]["data"]),
    #             dtype=getattr(numpy, loaded[k]["dtype"]),
    #         ).reshape(loaded[k]["shape"])
    #     )
    #     sf_tensor = reloaded[k]
    #     if not torch.equal(pt_tensor, sf_tensor):
    #         raise RuntimeError(f"The output tensors do not match for key {k}")


if __name__ == "__main__":
    try:
        convert_file(
            args.input,
            args.output,
            rename={
                "time_faaaa": "time_first",
                "time_maa": "time_mix",
                "lora_A": "lora.0",
                "lora_B": "lora.1",
            },
            transpose_names=[
                "time_mix_w1",
                "time_mix_w2",
                "time_decay_w1",
                "time_decay_w2",
                "w1",
                "w2",
                "a1",
                "a2",
                "g1",
                "g2",
                "v1",
                "v2",
                "time_state",
                "lora.0",
            ],
        )
        print(f"Saved to {args.output}")
    except Exception as e:
        print(e)
        with open("error.txt", "w") as f:
            f.write(str(e))