web-rwkv 0.1.1

An implementation of the RWKV language model in pure WebGPU.
Documentation 
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

import json
import os
import sys
import copy
import torch
from safetensors.torch import load_file, save_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 convert_file(
    pt_filename: str,
    sf_filename: str,
):
    loaded = torch.load(pt_filename, map_location="cpu")
    if "state_dict" in loaded:
        loaded = loaded["state_dict"]

    loaded = {k: v.clone().half() for k, v in loaded.items()}
    for k, v in loaded.items():
        print(f'{k}\t{v.shape}\t{v.dtype}')

    # For tensors to be contiguous
    loaded = {k: v.contiguous() for k, v in loaded.items()}

    dirname = os.path.dirname(sf_filename)
    os.makedirs(dirname, exist_ok=True)
    save_file(loaded, sf_filename, metadata={"format": "pt"})
    reloaded = load_file(sf_filename)
    for k in loaded:
        pt_tensor = loaded[k]
        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__":
    convert_file(args.input, args.output)
    print(f"Saved to {args.output}")