web-rwkv 0.8.13

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

#!/usr/bin/python

import json

import argparse

parser = argparse.ArgumentParser()
parser.add_argument(
    "--input",
    type=str,
    default="rwkv_vocab_v20230424.txt",
    help="Path to input txt")
parser.add_argument(
    "--output",
    type=str,
    default="rwkv_vocab_v20230424.json",
    help="Path to output JSON",
)
args = parser.parse_args()


I_TO_TOKEN = {}
lines = open(args.input, "r", encoding="utf-8").readlines()
for l in lines:
    idx = int(l[:l.index(' ')])
    x = eval(l[l.index(' '):l.rindex(' ')])
    if not isinstance(x, str):
        x = list(x)
    I_TO_TOKEN[idx] = x

out = open(args.output, "w")
out.write(json.dumps(I_TO_TOKEN, indent=4))