yum-slop/TaSTT

Free self-hosted STT for VRChat.

git clone https://git.yummers.dev/yum-slop/TaSTT

yumImport FastTextPager repof6b93a2

master
branchesmasterall refs →
4.2 KiB119 linesraw
1from math import ceil, floor
2from PIL import Image
3from unidecode import unidecode
4import sentencepiece as spm
5
6IMG_RES = 512  # square image
7
8def get_tokenizer():
9    use_sentencepiece = True
10
11    if not use_sentencepiece:
12        from tokenizers import Tokenizer
13        tokenizer_json = "./custom_wordpiece_tokenizer_65k/tokenizer.json"
14        print(f"Loading Tokenizers library tokenizer from: {tokenizer_json}")
15        return Tokenizer.from_file(tokenizer_json)
16    else:
17        model_path = "./custom_unigram_tokenizer_65k/unigram.model"
18        print(f"Loading SentencePiece tokenizer from: {model_path}")
19        sp = spm.SentencePieceProcessor()
20        sp.load(model_path)
21        print(f"Successfully loaded SentencePiece model. Vocab size: {sp.get_piece_size()}")
22        return sp
23
24def get_words():
25    tokenizer = get_tokenizer()
26
27    print(f"vocabulary size: {tokenizer.get_piece_size()}")
28    # sp_space = sentencepiece space.
29    # A special character sentencepiece uses to represent spaces before words.
30    sp_space = chr(9601)
31    words = []
32
33    # Accumulate words into a list, indexed by the token number. Sanitize them as
34    # you go.
35    for i in range(tokenizer.get_piece_size()):
36        word = tokenizer.id_to_piece(i)
37        tok = i
38        #print(f"  Original token ({tok}): {repr(word)} ({' '.join(str(ord(c)) for c in word)})")
39        word_sanitized = ""
40        # Dirty hack: convert non-ASCII characters to nearest ASCII equivalent
41        for c in word:
42            if ord(c) > 127 and c != sp_space:
43                c_plain = unidecode(c)
44                print(f"  Resolved {c} to {c_plain}")
45                word_sanitized += c_plain
46            else:
47                word_sanitized += c
48        # Replace sp_space with ' '
49        word_sanitized = word_sanitized.replace(sp_space, ' ')
50        #print(f"  {tok}: {word_sanitized}")
51        words.append(word_sanitized)
52
53    # Special word: empty string. SentencePiece doesn't support this natively.
54    words.append('')
55
56    return words
57
58# Fold a flat index into a IMG_RESxIMG_RES box. Return the (x,y) coordinate of
59# the folded index.
60def fold_idx(flat_idx):
61    return (flat_idx % IMG_RES, int(floor(flat_idx / IMG_RES)))
62
63def unfold_idx(coord):
64    return coord[0] + coord[1] * IMG_RES
65
66assert unfold_idx(fold_idx(1533125)) == 1533125
67assert unfold_idx(fold_idx(8538235)) == 8538235
68assert fold_idx(unfold_idx((192,235))) == (192,235)
69assert fold_idx(unfold_idx((83,388))) == (83,388)
70
71def generate_lut(words, filename):
72    # Write the texture header.
73    black = (0, 0, 0, 255)
74    img = Image.new('RGBA', (IMG_RES, IMG_RES), black)
75
76    # The header is `len(words)` slots long. Thus the actual LUT content starts at
77    # the index `len(words)`.
78    pixel_data = img.load()
79    lut_ptr = len(words)
80    for i in range(0, len(words)):
81        # Get pointer to the actual word data.
82        tok_ptr = lut_ptr
83        tok_len = len(words[i])
84        rgba = ((tok_ptr >>  0) & 0xFF,
85                (tok_ptr >>  8) & 0xFF,
86                (tok_ptr >> 16) & 0xFF,
87                tok_len)
88        print(f"Writing {rgba} to {i} / {fold_idx(i)}")
89        idx_x, idx_y = fold_idx(i)
90        pixel_data[idx_x, idx_y] = rgba
91
92        for j in range(0, ceil(tok_len/4.0)):
93            quad_ptr = tok_ptr + j
94            tok_0 = ord(words[i][j*4])
95            tok_1 = ord(words[i][j*4+1] if tok_len > j*4+1 else ' ')
96            tok_2 = ord(words[i][j*4+2] if tok_len > j*4+2 else ' ')
97            tok_3 = ord(words[i][j*4+3] if tok_len > j*4+3 else ' ')
98            rgba = (tok_0, tok_1, tok_2, tok_3)
99            idx_x, idx_y = fold_idx(quad_ptr)
100            print(f"  Writing {rgba} to {quad_ptr} / {fold_idx(quad_ptr)}")
101            pixel_data[idx_x, idx_y] = rgba
102
103        # Advance the LUT ptr. Since we store 4 chars per pixel (RGBA), we advance
104        # it by ceil(tok_len/4).
105        lut_ptr += int(ceil(tok_len/4.0))
106
107    pretty = False
108    if pretty:
109        for y in range(0, IMG_RES):
110            for x in range(0, IMG_RES):
111                rgba = pixel_data[x, y]
112                pixel_data[x, y] = (rgba[0], rgba[1], rgba[2], 255)
113
114    print(f"Saving to {filename}")
115    img.save(filename)
116
117if __name__ == "__main__":
118    words = get_words()
119    generate_lut(words, "bpe_lut.png")