yum/FastTextPager

Compressed text paging over OSC.

git clone https://git.yummers.dev/yum/FastTextPager

yumbugfixes790c91d

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