yum-slop/TaSTT

Free self-hosted STT for VRChat.

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

yumExperiment with hallucination reductiona7f9b7b

master
branchesmasterall refs →
21.1 KiB573 linesraw
1import app_config
2import argparse
3import io
4import keybind_event_machine
5from logger import log, log_err
6from math import floor, ceil
7import msvcrt
8import os
9import pygame
10from pythonosc import udp_client
11import sentencepiece as spm
12import steamvr
13from shared_thread_data import SharedThreadData
14import stt
15import sys
16import threading
17import time
18
19# Initialize pygame mixer
20pygame.mixer.init()
21
22TESTS_ENABLED = True
23
24# 0 = quiet, 1 = verbose, 2 = very verbose
25LOG_LEVEL = 0
26
27APP_ROOT = os.path.dirname(os.path.abspath(__file__))
28PROJECT_ROOT = os.path.dirname(APP_ROOT)
29
30def get_tokenizer():
31    model_path = os.path.join(PROJECT_ROOT, "custom_unigram_tokenizer_65k", "unigram.model")
32    log(f"Loading SentencePiece tokenizer from: {model_path}")
33    sp = spm.SentencePieceProcessor()
34    sp.load(model_path)
35    log(f"Successfully loaded SentencePiece model. Vocab size: {sp.get_piece_size()}")
36    return sp
37
38def parse_args():
39    parser = argparse.ArgumentParser()
40    parser.add_argument("--config", type=str, help="Path to config file (YAML).", required=True)
41    return parser.parse_args()
42
43def assert_equal(a, b):
44    err_msg = f"{a} != {b}"
45    assert a == b, err_msg
46
47# Turn a whitespace-delimited string into a list of strings no longer than
48# `cols`.
49# Preferentially breaks strings at whitespace boundaries. Preserves whitespace
50# between words, except if that whitespace comes between lines. Breaks words
51# longer than `cols` with a hyphen.
52def wrap_line(line: str, cols):
53    # First, split line into alternating chunks of words and whitespace.
54    def get_sequences(line):
55        is_space = False
56        sequences = []
57        seq_start = 0
58        seq_end = -1
59        for i in range(0, len(line)):
60            if line[i].isspace():
61                if is_space:
62                    seq_end = i
63                    continue
64                # We were looking at text, now we see whitespace.
65                seq = line[seq_start:seq_end+1]
66                if len(seq) > 0:
67                    sequences.append(seq)
68                seq_start = i
69                seq_end = i
70                is_space = True
71            else:
72                if not is_space:
73                    seq_end = i
74                    continue
75                # We were looking at whitespace, now we see text.
76                seq = line[seq_start:seq_end+1]
77                if len(seq) > 0:
78                    sequences.append(seq)
79                seq_start = i
80                seq_end = i
81                is_space = False
82        sequences.append(line[seq_start:seq_end+1])
83        return sequences
84    if TESTS_ENABLED:
85        assert_equal(get_sequences("foo"), ["foo"])
86        assert_equal(get_sequences("foo bar"), ["foo", " ", "bar"])
87        assert_equal(get_sequences(" foo bar"), [" ", "foo", " ", "bar"])
88        assert_equal(get_sequences(" foo  bar"), [" ", "foo", "  ", "bar"])
89
90    # Next, greedily construct lines out of those sequences.
91    # Whitespace gets treated specially. If it would push us over the limit, we
92    # end the line and drop the whitespace.
93    sequences = get_sequences(line)
94    def coalesce_sequences(sequences, cols):
95        cur_line = ""
96        lines = []
97        for seq in sequences:
98            if len(cur_line) + len(seq) <= cols:
99                cur_line += seq
100                continue
101            if seq.isspace():
102                lines.append(cur_line)
103                cur_line = ""
104                continue
105            if len(cur_line) > 0:
106                lines.append(cur_line)
107            # Edge case: text sequence is longer than a line.
108            while len(seq) > cols:
109                seq_prefix = seq[0:cols-1] + "-"
110                seq = seq[cols-1:]
111                lines.append(seq_prefix)
112            cur_line = seq
113        if len(cur_line) > 0:
114            lines.append(cur_line)
115        return lines
116    if TESTS_ENABLED:
117        assert_equal(coalesce_sequences(get_sequences("foo bar"), 3), ["foo", "bar"])
118        assert_equal(coalesce_sequences(get_sequences("foo bar"), 4), ["foo ", "bar"])
119        assert_equal(coalesce_sequences(get_sequences("foo  bar"), 4), ["foo", "bar"])
120        assert_equal(coalesce_sequences(get_sequences("foobar"), 3), ["fo-", "ob-", "ar"])
121        assert_equal(coalesce_sequences(get_sequences("f obar"), 3), ["f ", "ob-", "ar"])
122
123    lines = coalesce_sequences(sequences, cols)
124
125    # Next, pad each line with whitespace.
126    def pad_lines(lines, cols):
127        for i in range(0, len(lines)):
128            lines[i] += ' ' * (cols - len(lines[i]))
129        return lines
130    if TESTS_ENABLED:
131        assert_equal(pad_lines(["foo", "ba"], 4), ["foo ", "ba  "])
132        assert_equal(pad_lines(["foo"], 2), ["foo"])
133
134    return pad_lines(lines, cols)
135
136def get_blocks(lines, tokenizer, block_width, num_blocks):
137    if LOG_LEVEL == 2:
138        log(f"Lines sent to tokenizer: {''.join(lines)}")
139    tokens = tokenizer.encode_as_ids(''.join(lines))
140    if LOG_LEVEL == 2:
141        log(f"Tokens: {tokens}")
142    pieces = []
143    for tok in tokens:
144        piece = tokenizer.id_to_piece(tok)
145        pieces.append(piece)
146    if LOG_LEVEL == 2:
147        log(f"Pieces: {pieces}")
148
149    # Group tokens into blocks and pad with empty characters.
150    # Also get visual pointers - the location where each block will be rendered.
151    def get_blocks():
152        blocks = []
153        visual_pointer = 0
154        visual_pointers = []
155        for i in range(0, ceil(len(tokens) / block_width)):
156            visual_pointers.append(visual_pointer)
157            block = []
158            for j in range(0, block_width):
159                if i*block_width + j >= len(tokens):
160                    # Pad block with empty characters. 65535 is a special token.
161                    block += [65535] * (block_width - len(block))
162                    break
163                block.append(tokens[i*block_width+j])
164                visual_pointer += len(pieces[i*block_width+j])
165            blocks.append(block)
166        return (blocks, visual_pointers)
167    blocks, visual_pointers = get_blocks()
168    if LOG_LEVEL == 2:
169        log(f"Blocks: {blocks}")
170        log(f"Visual pointers: {visual_pointers}")
171
172    # Set all blocks up to the next `num_blocks` boundary to blank tokens.
173    # This handles the edge case where a prior message wrote data there which
174    # is covering up our new data.
175    def pad_blocks(blocks, visual_pointers):
176        cur_num_blocks = len(blocks)
177        num_pad_blocks = num_blocks - cur_num_blocks
178        for i in range(0, num_pad_blocks):
179            blocks.append([65535] * block_width)
180            visual_pointers.append(255)
181        return blocks, visual_pointers
182    blocks, visual_pointers = pad_blocks(blocks, visual_pointers)
183    if LOG_LEVEL == 2:
184        log(f"Blocks (padded): {blocks}")
185        log(f"Visual pointers (padded): {visual_pointers}")
186
187    return blocks, visual_pointers
188
189def calc_diff(prev_blocks, prev_visual_pointers, cur_blocks,
190              cur_visual_pointers):
191    diff_indices = []
192    diff_blocks = []
193    diff_visual_pointers = []
194
195    for i in range(0, len(cur_blocks)):
196        if i >= len(prev_blocks):
197            diff_blocks.append(cur_blocks[i])
198            diff_visual_pointers.append(cur_visual_pointers[i])
199            diff_indices.append(i)
200            continue
201        if prev_blocks[i] != cur_blocks[i] or prev_visual_pointers[i] != cur_visual_pointers[i]:
202            diff_blocks.append(cur_blocks[i])
203            diff_visual_pointers.append(cur_visual_pointers[i])
204            diff_indices.append(i)
205
206    return diff_indices, diff_blocks, diff_visual_pointers
207
208def send_data(osc_client, indices, blocks, visual_pointers):
209    def split_blocks_by_byte(blocks):
210        blocks_byte00 = []
211        blocks_byte01 = []
212        for block in blocks:
213            block_byte00 = []
214            block_byte01 = []
215            for datum in block:
216                block_byte00.append((datum >> 0) & 0xFF)
217                block_byte01.append((datum >> 8) & 0xFF)
218            blocks_byte00.append(block_byte00)
219            blocks_byte01.append(block_byte01)
220        return blocks_byte00, blocks_byte01
221
222    blocks_byte00, blocks_byte01 = split_blocks_by_byte(blocks)
223    if LOG_LEVEL == 2:
224        log(f"Blocks (byte 00): {blocks_byte00}")
225        log(f"Blocks (byte 01): {blocks_byte01}")
226
227    def send_osc(osc_client, addr, data):
228        osc_client.send_message(addr, data)
229
230    for i in range(0, len(blocks)):
231        lp_int = indices[i]
232        lp_param = "_Unigram_Letter_Grid_OSC_Pointer"
233        addr = "/avatar/parameters/" + lp_param
234        send_osc(osc_client, addr, lp_int)
235
236        vp_float = (-127.5 + visual_pointers[i]) / 127.5
237        vp_param = f"_Unigram_Letter_Grid_OSC_Visual_Pointer"
238        addr = "/avatar/parameters/" + vp_param
239        send_osc(osc_client, addr, vp_float)
240        if LOG_LEVEL == 2:
241            log(f"Sending block {blocks[i]} at {visual_pointers[i]} index {indices[i]}")
242        for j in range(0, len(blocks[i])):
243            byte00_float = (-127.5 + blocks_byte00[i][j]) / 127.5
244            byte01_float = (-127.5 + blocks_byte01[i][j]) / 127.5
245            byte00_param  = f"_Unigram_Letter_Grid_OSC_Datum{j:02}_Byte00"
246            byte01_param  = f"_Unigram_Letter_Grid_OSC_Datum{j:02}_Byte01"
247            addr = "/avatar/parameters/" + byte00_param
248            send_osc(osc_client, addr, byte00_float)
249            addr = "/avatar/parameters/" + byte01_param
250            send_osc(osc_client, addr, byte01_float)
251        time.sleep(0.34)
252
253def getOscClient(ip = "127.0.0.1", port = 9000):
254    return udp_client.SimpleUDPClient(ip, port)
255
256class InputState:
257    def __init__(self):
258        self.page = 0
259        # Initialize the known state of the board to empty array. This will cause
260        # our paging logic to re-send everything the first time around.
261        self.blocks = []
262        self.visual_pointers = []
263        pass
264
265def handle_input(state: InputState, line: str, tokenizer, osc_client, cfg):
266    line_wrapped = wrap_line(line, cfg["cols"])
267    if TESTS_ENABLED:
268        for line in line_wrapped:
269            assert_equal(len(line), cfg["cols"])
270    if LOG_LEVEL == 2:
271        log(f"Wrapped lines: {line_wrapped}")
272
273    # Get several blank lines whenever we roll over.
274    # It's better for the reader to have some continuity when the board pages
275    # over. If we simply replaced the entire screen, it would be harder to
276    # understand.
277    line_rollover = cfg["rows"] - 2
278    blank_line = ' ' * cfg["cols"]
279    # We show a full page, then only `line_rollover` additional lines per page.
280    end_ptr = cfg["rows"]
281    which_page = 0
282    while end_ptr < len(line_wrapped):
283        end_ptr += line_rollover
284        which_page += 1
285    if state.page != which_page:
286        state.blocks = []
287        state.visual_pointers = []
288        state.page = which_page
289    line_wrapped = line_wrapped[end_ptr-cfg["rows"]:]
290
291    # Get blocks and visual pointers.
292    blocks, visual_pointers = get_blocks(line_wrapped, tokenizer,
293                                         cfg["block_width"], cfg["num_blocks"])
294
295    # Note that because we only send one page of data at a time, we don't have
296    # to worry about wrapping visual pointers! We will basically never run out
297    # of space.
298    indices, diff_blocks, diff_visual_pointers = calc_diff(state.blocks, state.visual_pointers, blocks, visual_pointers)
299    indices = [idx % cfg["num_blocks"] for idx in indices]
300    # Send only one block at a time to make things snappier in interactive use
301    # case.
302    # TODO use a continuation (yield) instead of returning. Then we can be a
303    # little lighter on the cpu. Measurements show that this script is
304    # already very light but we're clearly wasting a lot of work by
305    # re-tokenizing the entire input every time we send a block.
306    if len(indices) == 0:
307        return
308    if indices[0] == len(state.blocks):
309        state.blocks.append(diff_blocks[0])
310        state.visual_pointers.append(diff_visual_pointers[0])
311    elif indices[0] > len(state.blocks):
312        log(f"This should never happen!")
313        sys.exit(1)
314    else:
315        state.blocks[indices[0]] = diff_blocks[0]
316        state.visual_pointers[indices[0]] = diff_visual_pointers[0]
317
318    send_data(osc_client, [indices[0]], [diff_blocks[0]], [diff_visual_pointers[0]])
319
320def osc_thread(shared_data: SharedThreadData):
321    osc_client = getOscClient()
322
323    def join_segments(a, b):
324        if len(a) > 0 and a[-1] != ' ':
325            return a + ' ' + b
326        else:
327            return a + b
328
329    if shared_data.cfg["use_builtin"]:
330        last_change = time.time()
331        remote_word = ""
332        while not shared_data.exit_event.is_set():
333            time.sleep(0.1)
334            local_word = ""
335            with shared_data.word_lock:
336                local_word = join_segments(shared_data.transcript,
337                                           shared_data.preview)
338            local_word = local_word[-140:]
339            if local_word == remote_word:
340                continue
341            if time.time() - last_change < 1.5:
342                continue
343            addr = "/chatbox/input"
344            if shared_data.cfg["enable_debug_mode"]:
345                log(f"Send {local_word}")
346            osc_client.send_message(addr, (local_word, True, False))
347            last_change = time.time()
348            remote_word = local_word
349    else:
350        # Custom chatbox
351        tokenizer = get_tokenizer()
352
353        # Prime the board
354        log("Priming the board")
355        input_state = InputState()
356        handle_input(input_state, "", tokenizer, osc_client, shared_data.cfg)
357
358        while not shared_data.exit_event.is_set():
359            word_copy = ""
360            with shared_data.word_lock:
361                word_copy = shared_data.word
362            handle_input(input_state, word_copy, tokenizer, osc_client, shared_data.cfg)
363            time.sleep(0.01)
364
365
366def vrInputThread(shared_data: SharedThreadData):
367    RECORD_STATE = 0
368    PAUSE_STATE = 1
369    state = PAUSE_STATE
370
371    hand_id = shared_data.cfg["button_hand"]
372    button_id = shared_data.cfg["button_type"]
373
374    # Rough description of state machine:
375    #   Single short press: toggle transcription
376    #   Medium press: dismiss custom chatbox
377    #   Long press: update chatbox in place
378    #   Medium press + long press: type transcription
379
380    last_rising = time.time()
381    last_medium_press_end = 0
382
383    waveform0 =  os.path.join(PROJECT_ROOT, "Sounds/Noise_On_Quiet.wav")
384    waveform1 =  os.path.join(PROJECT_ROOT, "Sounds/Noise_Off_Quiet.wav")
385    waveform2 =  os.path.join(PROJECT_ROOT, "Sounds/Dismiss_Noise_Quiet.wav")
386    waveform3 =  os.path.join(PROJECT_ROOT, "Sounds/KB_Noise_Off_Quiet.wav")
387
388    button_generator = steamvr.pollButtonPress(hand=hand_id, button=button_id,
389            shared_data=shared_data)
390    while not shared_data.exit_event.is_set():
391        time.sleep(0.01)
392        try:
393            event = next(button_generator)
394        except StopIteration:
395            break
396
397        with shared_data.word_lock:
398            if not shared_data.stream or not shared_data.collector:
399                continue
400
401            if event.opcode == steamvr.EVENT_RISING_EDGE:
402                last_rising = time.time()
403
404                if state == PAUSE_STATE:
405                    shared_data.stream.pause(False)
406                    shared_data.stream.getSamples()
407
408            elif event.opcode == steamvr.EVENT_FALLING_EDGE:
409                now = time.time()
410                if now - last_rising > 1.5:
411                    # Long press: treat as the end of transcription.
412                    state = PAUSE_STATE
413
414                    shared_data.stream.pause(True)
415
416                    if last_rising - last_medium_press_end < 1.0:
417                        # Type transcription
418                        play_sound_with_volume(waveform3, shared_data.cfg)
419                    else:
420                        play_sound_with_volume(waveform1, shared_data.cfg)
421
422                elif now - last_rising > 0.5:
423                    # Medium press
424                    log_err("CLEARING")
425                    last_medium_press_end = now
426                    state = PAUSE_STATE
427                    play_sound_with_volume(waveform2, shared_data.cfg)
428
429                    # Flush the *entire* pipeline.
430                    shared_data.stream.pause(True)
431                    shared_data.stream.getSamples()
432                    shared_data.collector.dropAudio()
433                    shared_data.transcript = ""
434                    shared_data.preview = ""
435                    continue
436
437                # Short hold
438                if state == RECORD_STATE:
439                    log_err("PAUSED")
440                    state = PAUSE_STATE
441
442                    shared_data.stream.pause(True)
443                    play_sound_with_volume(waveform1, shared_data.cfg)
444                elif state == PAUSE_STATE:
445                    log_err("RECORDING")
446                    state = RECORD_STATE
447                    if shared_data.cfg["reset_on_toggle"]:
448                        if shared_data.cfg["enable_debug_mode"]:
449                            log_err("Toggle detected, dropping transcript (3)")
450                        shared_data.transcript = ""
451                        shared_data.preview = ""
452                        #audio_state.drop_transcription = True
453                    else:
454                        if shared_data.cfg["enable_debug_mode"]:
455                            log_err("Toggle detected, committing preview text (3)")
456                        #audio_state.text += audio_state.preview_text
457
458                    shared_data.stream.pause(False)
459                    play_sound_with_volume(waveform0, shared_data.cfg)
460
461
462def kbInputThread(shared_data: SharedThreadData):
463    machine = keybind_event_machine.KeybindEventMachine(shared_data.cfg["keybind"])
464    last_press_time = 0
465
466    # double pressing the keybind
467    double_press_timeout = 0.5
468
469    RECORD_STATE = 0
470    PAUSE_STATE = 1
471    state = PAUSE_STATE
472
473    waveform0 = os.path.join(PROJECT_ROOT, "Sounds/Noise_On_Quiet.wav")
474    waveform1 = os.path.join(PROJECT_ROOT, "Sounds/Noise_Off_Quiet.wav")
475    waveform2 = os.path.join(PROJECT_ROOT, "Sounds/Dismiss_Noise_Quiet.wav")
476    waveform3 = os.path.join(PROJECT_ROOT, "Sounds/KB_Noise_Off_Quiet.wav")
477
478    while not shared_data.exit_event.is_set():
479        time.sleep(0.01)
480
481        cur_press_time = machine.getNextPressTime()
482        if cur_press_time == 0:
483            continue
484
485        with shared_data.word_lock:
486            if not shared_data.stream or not shared_data.collector:
487                continue
488
489            EVENT_SINGLE_PRESS = 0
490            EVENT_DOUBLE_PRESS = 1
491            if last_press_time == 0:
492                event = EVENT_SINGLE_PRESS
493            elif cur_press_time - last_press_time < double_press_timeout:
494                event = EVENT_DOUBLE_PRESS
495            else:
496                event = EVENT_SINGLE_PRESS
497            last_press_time = cur_press_time
498
499            if event == EVENT_DOUBLE_PRESS:
500                log_err("CLEARING")
501                state = PAUSE_STATE
502                play_sound_with_volume(waveform2, shared_data.cfg)
503
504                # Flush the *entire* pipeline.
505                shared_data.stream.pause(True)
506                shared_data.stream.getSamples()
507                shared_data.collector.dropAudio()
508                shared_data.transcript = ""
509                shared_data.preview = ""
510                continue
511
512            # Short hold
513            if state == RECORD_STATE:
514                log_err("PAUSED")
515                state = PAUSE_STATE
516                shared_data.stream.pause(True)
517                play_sound_with_volume(waveform1, shared_data.cfg)
518            elif state == PAUSE_STATE:
519                log_err("RECORDING")
520                state = RECORD_STATE
521                if shared_data.cfg["reset_on_toggle"]:
522                    if shared_data.cfg["enable_debug_mode"]:
523                        log_err("Toggle detected, dropping transcript (2)")
524                    shared_data.transcript = ""
525                    shared_data.preview = ""
526                else:
527                    if shared_data.cfg["enable_debug_mode"]:
528                        log_err("Toggle detected, committing preview text (2)")
529                shared_data.stream.pause(False)
530                play_sound_with_volume(waveform0, shared_data.cfg)
531
532def play_sound_with_volume(filepath, cfg):
533    """Play a WAV file with adjusted volume"""
534    volume = cfg.get("volume", 30)
535    
536    try:
537        sound = pygame.mixer.Sound(filepath)
538        sound.set_volume(volume * 0.01)
539        sound.play()
540    except Exception as e:
541        log_err(f"Error playing sound {filepath}: {e}")
542
543if __name__ == "__main__":
544    cli_args = parse_args()
545    cfg = app_config.getConfig(cli_args.config)
546    shared_data = SharedThreadData(cfg)
547    osc_thread = threading.Thread(
548            target=osc_thread,
549            args=(shared_data,))
550    osc_thread.start()
551
552    transcribe_thread = threading.Thread(
553            target=stt.transcriptionThread,
554            args=(shared_data,))
555    transcribe_thread.start()
556
557    vr_input_thd = threading.Thread(target=vrInputThread, args=(shared_data,))
558    vr_input_thd.start()
559
560    kb_input_thd = threading.Thread(target=kbInputThread, args=(shared_data,))
561    kb_input_thd.start()
562
563    word_is_over = False
564    local_word = ""
565    while True:
566        time.sleep(0.1)
567        continue
568    shared_data.exit_event.set()
569    osc_thread.join()
570    transcribe_thread.join()
571    vr_input_thd.join()
572    kb_input_thd.join()
573