diff options
Diffstat (limited to 'app')
| -rw-r--r-- | app/app_config.py | 39 | ||||
| -rw-r--r-- | app/hi.py | 580 | ||||
| -rw-r--r-- | app/keybind_event_machine.py | 21 | ||||
| -rw-r--r-- | app/list_microphones.py | 24 | ||||
| -rw-r--r-- | app/profanity_filter.py | 43 | ||||
| -rw-r--r-- | app/requirements.txt | 12 | ||||
| -rw-r--r-- | app/shared_thread_data.py | 14 | ||||
| -rw-r--r-- | app/steamvr.py | 87 | ||||
| -rw-r--r-- | app/stt.py | 915 |
9 files changed, 1735 insertions, 0 deletions
diff --git a/app/app_config.py b/app/app_config.py new file mode 100644 index 0000000..f911456 --- /dev/null +++ b/app/app_config.py @@ -0,0 +1,39 @@ +import os +import sys +import typing + +def getConfig(path: str) -> typing.Dict[str, typing.Union[str, float, int, bool]]: + # Helper functions to detect and convert the type + def is_int(value: str) -> bool: + try: + int(value) + return True + except ValueError: + return False + + def is_float(value: str) -> bool: + try: + float(value) + return True + except ValueError: + return False + + def convert_value(key: str, value: str): + if key.startswith(("enable_", "remove_", "use_", "clear_")): + return bool(int(value)) + elif is_int(value): + return int(value) + elif is_float(value): + return float(value) + else: + return value + + config = {} + with open(path, 'r') as file: + for line in file: + key_value = line.strip().split(": ", maxsplit=1) + key = key_value[0] + value = key_value[1] if len(key_value) > 1 else "" + config[key] = convert_value(key, value.strip()) + return config + diff --git a/app/hi.py b/app/hi.py new file mode 100644 index 0000000..bb09418 --- /dev/null +++ b/app/hi.py @@ -0,0 +1,580 @@ +import app_config +import argparse +import io +import keybind_event_machine +from math import floor, ceil +import msvcrt +import os +from pythonosc import udp_client +import sentencepiece as spm +import steamvr +from shared_thread_data import SharedThreadData +import stt +import sys +import threading +import time +import pygame + +sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8') +sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding='utf-8') + +# Initialize pygame mixer +pygame.mixer.init() + +TESTS_ENABLED = True + +# 0 = quiet, 1 = verbose, 2 = very verbose +LOG_LEVEL = 0 + +APP_ROOT = os.path.dirname(os.path.abspath(__file__)) +PROJECT_ROOT = os.path.dirname(APP_ROOT) + +def get_tokenizer(): + model_path = os.path.join(PROJECT_ROOT, "custom_unigram_tokenizer_65k", "unigram.model") + print(f"Loading SentencePiece tokenizer from: {model_path}") + sp = spm.SentencePieceProcessor() + sp.load(model_path) + print(f"Successfully loaded SentencePiece model. Vocab size: {sp.get_piece_size()}") + return sp + +def parse_args(): + parser = argparse.ArgumentParser() + parser.add_argument("--config", type=str, help="Path to config file (YAML).", required=True) + return parser.parse_args() + +def assert_equal(a, b): + err_msg = f"{a} != {b}" + assert a == b, err_msg + +# Turn a whitespace-delimited string into a list of strings no longer than +# `cols`. +# Preferentially breaks strings at whitespace boundaries. Preserves whitespace +# between words, except if that whitespace comes between lines. Breaks words +# longer than `cols` with a hyphen. +def wrap_line(line: str, cols): + # First, split line into alternating chunks of words and whitespace. + def get_sequences(line): + is_space = False + sequences = [] + seq_start = 0 + seq_end = -1 + for i in range(0, len(line)): + if line[i].isspace(): + if is_space: + seq_end = i + continue + # We were looking at text, now we see whitespace. + seq = line[seq_start:seq_end+1] + if len(seq) > 0: + sequences.append(seq) + seq_start = i + seq_end = i + is_space = True + else: + if not is_space: + seq_end = i + continue + # We were looking at whitespace, now we see text. + seq = line[seq_start:seq_end+1] + if len(seq) > 0: + sequences.append(seq) + seq_start = i + seq_end = i + is_space = False + sequences.append(line[seq_start:seq_end+1]) + return sequences + if TESTS_ENABLED: + assert_equal(get_sequences("foo"), ["foo"]) + assert_equal(get_sequences("foo bar"), ["foo", " ", "bar"]) + assert_equal(get_sequences(" foo bar"), [" ", "foo", " ", "bar"]) + assert_equal(get_sequences(" foo bar"), [" ", "foo", " ", "bar"]) + + # Next, greedily construct lines out of those sequences. + # Whitespace gets treated specially. If it would push us over the limit, we + # end the line and drop the whitespace. + sequences = get_sequences(line) + def coalesce_sequences(sequences, cols): + cur_line = "" + lines = [] + for seq in sequences: + if len(cur_line) + len(seq) <= cols: + cur_line += seq + continue + if seq.isspace(): + lines.append(cur_line) + cur_line = "" + continue + if len(cur_line) > 0: + lines.append(cur_line) + # Edge case: text sequence is longer than a line. + while len(seq) > cols: + seq_prefix = seq[0:cols-1] + "-" + seq = seq[cols-1:] + lines.append(seq_prefix) + cur_line = seq + if len(cur_line) > 0: + lines.append(cur_line) + return lines + if TESTS_ENABLED: + assert_equal(coalesce_sequences(get_sequences("foo bar"), 3), ["foo", "bar"]) + assert_equal(coalesce_sequences(get_sequences("foo bar"), 4), ["foo ", "bar"]) + assert_equal(coalesce_sequences(get_sequences("foo bar"), 4), ["foo", "bar"]) + assert_equal(coalesce_sequences(get_sequences("foobar"), 3), ["fo-", "ob-", "ar"]) + assert_equal(coalesce_sequences(get_sequences("f obar"), 3), ["f ", "ob-", "ar"]) + + lines = coalesce_sequences(sequences, cols) + + # Next, pad each line with whitespace. + def pad_lines(lines, cols): + for i in range(0, len(lines)): + lines[i] += ' ' * (cols - len(lines[i])) + return lines + if TESTS_ENABLED: + assert_equal(pad_lines(["foo", "ba"], 4), ["foo ", "ba "]) + assert_equal(pad_lines(["foo"], 2), ["foo"]) + + return pad_lines(lines, cols) + +def get_blocks(lines, tokenizer, block_width, num_blocks): + if LOG_LEVEL == 2: + print(f"Lines sent to tokenizer: {''.join(lines)}") + tokens = tokenizer.encode_as_ids(''.join(lines)) + if LOG_LEVEL == 2: + print(f"Tokens: {tokens}") + pieces = [] + for tok in tokens: + piece = tokenizer.id_to_piece(tok) + pieces.append(piece) + if LOG_LEVEL == 2: + print(f"Pieces: {pieces}") + + # Group tokens into blocks and pad with empty characters. + # Also get visual pointers - the location where each block will be rendered. + def get_blocks(): + blocks = [] + visual_pointer = 0 + visual_pointers = [] + for i in range(0, ceil(len(tokens) / block_width)): + visual_pointers.append(visual_pointer) + block = [] + for j in range(0, block_width): + if i*block_width + j >= len(tokens): + # Pad block with empty characters. 65535 is a special token. + block += [65535] * (block_width - len(block)) + break + block.append(tokens[i*block_width+j]) + visual_pointer += len(pieces[i*block_width+j]) + blocks.append(block) + return (blocks, visual_pointers) + blocks, visual_pointers = get_blocks() + if LOG_LEVEL == 2: + print(f"Blocks: {blocks}") + print(f"Visual pointers: {visual_pointers}") + + # Set all blocks up to the next `num_blocks` boundary to blank tokens. + # This handles the edge case where a prior message wrote data there which + # is covering up our new data. + def pad_blocks(blocks, visual_pointers): + cur_num_blocks = len(blocks) + num_pad_blocks = num_blocks - cur_num_blocks + for i in range(0, num_pad_blocks): + blocks.append([65535] * block_width) + visual_pointers.append(255) + return blocks, visual_pointers + blocks, visual_pointers = pad_blocks(blocks, visual_pointers) + if LOG_LEVEL == 2: + print(f"Blocks (padded): {blocks}") + print(f"Visual pointers (padded): {visual_pointers}") + + return blocks, visual_pointers + +def calc_diff(prev_blocks, prev_visual_pointers, cur_blocks, + cur_visual_pointers): + diff_indices = [] + diff_blocks = [] + diff_visual_pointers = [] + + for i in range(0, len(cur_blocks)): + if i >= len(prev_blocks): + diff_blocks.append(cur_blocks[i]) + diff_visual_pointers.append(cur_visual_pointers[i]) + diff_indices.append(i) + continue + if prev_blocks[i] != cur_blocks[i] or prev_visual_pointers[i] != cur_visual_pointers[i]: + diff_blocks.append(cur_blocks[i]) + diff_visual_pointers.append(cur_visual_pointers[i]) + diff_indices.append(i) + + return diff_indices, diff_blocks, diff_visual_pointers + +def send_data(osc_client, indices, blocks, visual_pointers): + def split_blocks_by_byte(blocks): + blocks_byte00 = [] + blocks_byte01 = [] + for block in blocks: + block_byte00 = [] + block_byte01 = [] + for datum in block: + block_byte00.append((datum >> 0) & 0xFF) + block_byte01.append((datum >> 8) & 0xFF) + blocks_byte00.append(block_byte00) + blocks_byte01.append(block_byte01) + return blocks_byte00, blocks_byte01 + + blocks_byte00, blocks_byte01 = split_blocks_by_byte(blocks) + if LOG_LEVEL == 2: + print(f"Blocks (byte 00): {blocks_byte00}") + print(f"Blocks (byte 01): {blocks_byte01}") + + def send_osc(osc_client, addr, data): + #print(f"Sending {data} to {addr}") + osc_client.send_message(addr, data) + + for i in range(0, len(blocks)): + lp_int = indices[i] + lp_param = "_Unigram_Letter_Grid_OSC_Pointer" + addr = "/avatar/parameters/" + lp_param + send_osc(osc_client, addr, lp_int) + + vp_float = (-127.5 + visual_pointers[i]) / 127.5 + vp_param = f"_Unigram_Letter_Grid_OSC_Visual_Pointer" + addr = "/avatar/parameters/" + vp_param + send_osc(osc_client, addr, vp_float) + if LOG_LEVEL == 2: + print(f"Sending block {blocks[i]} at {visual_pointers[i]} index {indices[i]}") + for j in range(0, len(blocks[i])): + byte00_float = (-127.5 + blocks_byte00[i][j]) / 127.5 + byte01_float = (-127.5 + blocks_byte01[i][j]) / 127.5 + byte00_param = f"_Unigram_Letter_Grid_OSC_Datum{j:02}_Byte00" + byte01_param = f"_Unigram_Letter_Grid_OSC_Datum{j:02}_Byte01" + addr = "/avatar/parameters/" + byte00_param + send_osc(osc_client, addr, byte00_float) + addr = "/avatar/parameters/" + byte01_param + send_osc(osc_client, addr, byte01_float) + time.sleep(0.34) + +def getOscClient(ip = "127.0.0.1", port = 9000): + return udp_client.SimpleUDPClient(ip, port) + +class InputState: + def __init__(self): + self.page = 0 + # Initialize the known state of the board to empty array. This will cause + # our paging logic to re-send everything the first time around. + self.blocks = [] + self.visual_pointers = [] + pass + +def handle_input(state: InputState, line: str, tokenizer, osc_client, cfg): + line_wrapped = wrap_line(line, cfg["cols"]) + if TESTS_ENABLED: + for line in line_wrapped: + assert_equal(len(line), cfg["cols"]) + if LOG_LEVEL == 2: + print(f"Wrapped lines: {line_wrapped}") + + # Get several blank lines whenever we roll over. + # It's better for the reader to have some continuity when the board pages + # over. If we simply replaced the entire screen, it would be harder to + # understand. + line_rollover = cfg["rows"] - 2 + blank_line = ' ' * cfg["cols"] + # We show a full page, then only `line_rollover` additional lines per page. + end_ptr = cfg["rows"] + which_page = 0 + while end_ptr < len(line_wrapped): + end_ptr += line_rollover + which_page += 1 + if state.page != which_page: + state.blocks = [] + state.visual_pointers = [] + state.page = which_page + line_wrapped = line_wrapped[end_ptr-cfg["rows"]:] + + # Get blocks and visual pointers. + blocks, visual_pointers = get_blocks(line_wrapped, tokenizer, + cfg["block_width"], cfg["num_blocks"]) + + # Note that because we only send one page of data at a time, we don't have + # to worry about wrapping visual pointers! We will basically never run out + # of space. + indices, diff_blocks, diff_visual_pointers = calc_diff(state.blocks, state.visual_pointers, blocks, visual_pointers) + indices = [idx % cfg["num_blocks"] for idx in indices] + # Send only one block at a time to make things snappier in interactive use + # case. + # TODO use a continuation (yield) instead of returning. Then we can be a + # little lighter on the cpu. Measurements show that this script is + # already very light but we're clearly wasting a lot of work by + # re-tokenizing the entire input every time we send a block. + if len(indices) == 0: + return + if indices[0] == len(state.blocks): + state.blocks.append(diff_blocks[0]) + state.visual_pointers.append(diff_visual_pointers[0]) + elif indices[0] > len(state.blocks): + print(f"This should never happen!") + sys.exit(1) + else: + state.blocks[indices[0]] = diff_blocks[0] + state.visual_pointers[indices[0]] = diff_visual_pointers[0] + + send_data(osc_client, [indices[0]], [diff_blocks[0]], [diff_visual_pointers[0]]) + +def osc_thread(shared_data: SharedThreadData): + osc_client = getOscClient() + + def join_segments(a, b): + if len(a) > 0 and a[-1] != ' ': + return a + ' ' + b + else: + return a + b + + if shared_data.cfg["use_builtin"]: + last_change = time.time() + remote_word = "" + while not shared_data.exit_event.is_set(): + time.sleep(0.1) + local_word = "" + with shared_data.word_lock: + local_word = join_segments(shared_data.transcript, + shared_data.preview) + local_word = local_word[-140:] + if local_word == remote_word: + continue + if time.time() - last_change < 1.5: + continue + addr = "/chatbox/input" + if shared_data.cfg["enable_debug_mode"]: + print(f"Send {local_word}", flush=True) + osc_client.send_message(addr, (local_word, True, False)) + last_change = time.time() + remote_word = local_word + else: + # Custom chatbox + tokenizer = get_tokenizer() + + # Prime the board + print("Priming the board") + input_state = InputState() + handle_input(input_state, "", tokenizer, osc_client, shared_data.cfg) + + while not shared_data.exit_event.is_set(): + word_copy = "" + with shared_data.word_lock: + word_copy = shared_data.word + handle_input(input_state, word_copy, tokenizer, osc_client, shared_data.cfg) + time.sleep(0.01) + + +def vrInputThread(shared_data: SharedThreadData): + RECORD_STATE = 0 + PAUSE_STATE = 1 + state = PAUSE_STATE + + hand_id = shared_data.cfg["button_hand"] + button_id = shared_data.cfg["button_type"] + + # Rough description of state machine: + # Single short press: toggle transcription + # Medium press: dismiss custom chatbox + # Long press: update chatbox in place + # Medium press + long press: type transcription + + last_rising = time.time() + last_medium_press_end = 0 + + waveform0 = os.path.join(PROJECT_ROOT, "Sounds/Noise_On_Quiet.wav") + waveform1 = os.path.join(PROJECT_ROOT, "Sounds/Noise_Off_Quiet.wav") + waveform2 = os.path.join(PROJECT_ROOT, "Sounds/Dismiss_Noise_Quiet.wav") + waveform3 = os.path.join(PROJECT_ROOT, "Sounds/KB_Noise_Off_Quiet.wav") + + button_generator = steamvr.pollButtonPress(hand=hand_id, button=button_id, + shared_data=shared_data) + while not shared_data.exit_event.is_set(): + time.sleep(0.01) + try: + event = next(button_generator) + except StopIteration: + break + + with shared_data.word_lock: + if not shared_data.stream or not shared_data.collector: + continue + + if event.opcode == steamvr.EVENT_RISING_EDGE: + last_rising = time.time() + + if state == PAUSE_STATE: + shared_data.stream.pause(False) + shared_data.stream.getSamples() + + elif event.opcode == steamvr.EVENT_FALLING_EDGE: + now = time.time() + if now - last_rising > 1.5: + # Long press: treat as the end of transcription. + state = PAUSE_STATE + + shared_data.stream.pause(True) + + if last_rising - last_medium_press_end < 1.0: + # Type transcription + play_sound_with_volume(waveform3, shared_data.cfg) + else: + play_sound_with_volume(waveform1, shared_data.cfg) + + elif now - last_rising > 0.5: + # Medium press + print("CLEARING", file=sys.stderr) + last_medium_press_end = now + state = PAUSE_STATE + play_sound_with_volume(waveform2, shared_data.cfg) + + # Flush the *entire* pipeline. + shared_data.stream.pause(True) + shared_data.stream.getSamples() + shared_data.collector.dropAudio() + shared_data.transcript = "" + shared_data.preview = "" + continue + + # Short hold + if state == RECORD_STATE: + print("PAUSED", file=sys.stderr) + state = PAUSE_STATE + + shared_data.stream.pause(True) + play_sound_with_volume(waveform1, shared_data.cfg) + elif state == PAUSE_STATE: + print("RECORDING", file=sys.stderr) + state = RECORD_STATE + if shared_data.cfg["reset_on_toggle"]: + if shared_data.cfg["enable_debug_mode"]: + print("Toggle detected, dropping transcript (3)", + file=sys.stderr) + shared_data.transcript = "" + shared_data.preview = "" + #audio_state.drop_transcription = True + else: + if shared_data.cfg["enable_debug_mode"]: + print("Toggle detected, committing preview text (3)", + file=sys.stderr) + #audio_state.text += audio_state.preview_text + + shared_data.stream.pause(False) + play_sound_with_volume(waveform0, shared_data.cfg) + + +def kbInputThread(shared_data: SharedThreadData): + machine = keybind_event_machine.KeybindEventMachine(shared_data.cfg["keybind"]) + last_press_time = 0 + + # double pressing the keybind + double_press_timeout = 0.5 + + RECORD_STATE = 0 + PAUSE_STATE = 1 + state = PAUSE_STATE + + waveform0 = os.path.join(PROJECT_ROOT, "Sounds/Noise_On_Quiet.wav") + waveform1 = os.path.join(PROJECT_ROOT, "Sounds/Noise_Off_Quiet.wav") + waveform2 = os.path.join(PROJECT_ROOT, "Sounds/Dismiss_Noise_Quiet.wav") + waveform3 = os.path.join(PROJECT_ROOT, "Sounds/KB_Noise_Off_Quiet.wav") + + while not shared_data.exit_event.is_set(): + time.sleep(0.01) + + cur_press_time = machine.getNextPressTime() + if cur_press_time == 0: + continue + + with shared_data.word_lock: + if not shared_data.stream or not shared_data.collector: + continue + + EVENT_SINGLE_PRESS = 0 + EVENT_DOUBLE_PRESS = 1 + if last_press_time == 0: + event = EVENT_SINGLE_PRESS + elif cur_press_time - last_press_time < double_press_timeout: + event = EVENT_DOUBLE_PRESS + else: + event = EVENT_SINGLE_PRESS + last_press_time = cur_press_time + + if event == EVENT_DOUBLE_PRESS: + print("CLEARING", file=sys.stderr) + state = PAUSE_STATE + play_sound_with_volume(waveform2, shared_data.cfg) + + # Flush the *entire* pipeline. + shared_data.stream.pause(True) + shared_data.stream.getSamples() + shared_data.collector.dropAudio() + shared_data.transcript = "" + shared_data.preview = "" + continue + + # Short hold + if state == RECORD_STATE: + print("PAUSED", file=sys.stderr) + state = PAUSE_STATE + shared_data.stream.pause(True) + play_sound_with_volume(waveform1, shared_data.cfg) + elif state == PAUSE_STATE: + print("RECORDING", file=sys.stderr) + state = RECORD_STATE + if shared_data.cfg["reset_on_toggle"]: + if shared_data.cfg["enable_debug_mode"]: + print("Toggle detected, dropping transcript (2)", + file=sys.stderr) + shared_data.transcript = "" + shared_data.preview = "" + else: + if shared_data.cfg["enable_debug_mode"]: + print("Toggle detected, committing preview text (2)", + file=sys.stderr) + shared_data.stream.pause(False) + play_sound_with_volume(waveform0, shared_data.cfg) + +def play_sound_with_volume(filepath, cfg): + """Play a WAV file with adjusted volume""" + volume = cfg.get("volume", 30) + + try: + sound = pygame.mixer.Sound(filepath) + sound.set_volume(volume * 0.01) + sound.play() + except Exception as e: + print(f"Error playing sound {filepath}: {e}", file=sys.stderr) + +if __name__ == "__main__": + cli_args = parse_args() + cfg = app_config.getConfig(cli_args.config) + shared_data = SharedThreadData(cfg) + osc_thread = threading.Thread( + target=osc_thread, + args=(shared_data,)) + osc_thread.start() + + transcribe_thread = threading.Thread( + target=stt.transcriptionThread, + args=(shared_data,)) + transcribe_thread.start() + + vr_input_thd = threading.Thread(target=vrInputThread, args=(shared_data,)) + vr_input_thd.start() + + kb_input_thd = threading.Thread(target=kbInputThread, args=(shared_data,)) + kb_input_thd.start() + + word_is_over = False + local_word = "" + while True: + time.sleep(0.1) + continue + shared_data.exit_event.set() + osc_thread.join() + transcribe_thread.join() + vr_input_thd.join() + kb_input_thd.join() + diff --git a/app/keybind_event_machine.py b/app/keybind_event_machine.py new file mode 100644 index 0000000..3ce6794 --- /dev/null +++ b/app/keybind_event_machine.py @@ -0,0 +1,21 @@ +import keyboard +import time + +class KeybindEventMachine: + def __init__(self, keybind: str): + self.keybind = keybind + self.events = [] + keyboard.add_hotkey(keybind, self.onPress) + + def onPress(self) -> None: + self.events.append(time.time()) + + # Returns the timestamp when the keybind was pressed, or 0 if no keypresses + # are queued. + def getNextPressTime(self) -> int: + if len(self.events) == 0: + return 0 + ret = self.events[0] + self.events = self.events[1:] + return ret + diff --git a/app/list_microphones.py b/app/list_microphones.py new file mode 100644 index 0000000..a6b1f36 --- /dev/null +++ b/app/list_microphones.py @@ -0,0 +1,24 @@ +import pyaudio +import json +import sys + +try: + p = pyaudio.PyAudio() + info = p.get_host_api_info_by_index(0) + numdevices = info.get('deviceCount') + + microphones = [] + for i in range(0, numdevices): + device_info = p.get_device_info_by_host_api_device_index(0, i) + if device_info.get('maxInputChannels') > 0: + microphones.append({ + 'index': i, + 'name': device_info.get('name'), + 'defaultSampleRate': device_info.get('defaultSampleRate') + }) + + print(json.dumps(microphones)) + p.terminate() +except Exception as e: + print(json.dumps({'error': str(e)}), file=sys.stderr) + sys.exit(1)
\ No newline at end of file diff --git a/app/profanity_filter.py b/app/profanity_filter.py new file mode 100644 index 0000000..b8c84ed --- /dev/null +++ b/app/profanity_filter.py @@ -0,0 +1,43 @@ +#!/usr/bin/env python3 + +class ProfanityFilter: + def __init__(self, en_path: str): + self.en_path = en_path + self.en_profanity = set() + + def load(self): + with open(self.en_path, 'r') as f: + for line in f: + self.en_profanity.add(line.strip()) + + def filter(self, line: str, language_code: str = "en") -> str: + filtered = "" + + if language_code not in {"en"}: + raise ValueError(f"Language code \"{language_code}\" is " + + "unsupported by the profanity filter") + + # Translation table converting vowels to asterisks. + vowel_to_asterisk = str.maketrans('aeiouAEIOU', '**********') + + result = [] + for word in line.split(): + word_clean = word.lower() + # Filter out non-alphabet characters from the word. + word_clean = ''.join([char for char in word_clean if char.isalpha()]) + if word_clean in self.en_profanity: + result.append(word.translate(vowel_to_asterisk)) + else: + result.append(word) + + return " ".join(result) + +if __name__ == "__main__": + en_path = "/mnt/d/vrc/TaSTT/GUI/Profanity/Profanity/en" + p = ProfanityFilter(en_path) + p.load() + assert(p.filter("fuck") == "f*ck") + assert(p.filter("fuck!") == "f*ck!") + assert(p.filter("fuck shit") == "f*ck sh*t") + assert(p.filter("fuck shit this should not be filtered") == "f*ck sh*t this should not be filtered") + assert(p.filter("ASS") == "*SS") diff --git a/app/requirements.txt b/app/requirements.txt new file mode 100644 index 0000000..c8d69df --- /dev/null +++ b/app/requirements.txt @@ -0,0 +1,12 @@ +faster-whisper +hf-xet +keyboard +langcodes +noisereduce +pyaudio +pygame +pydub +python-osc +sentencepiece +silero-vad +openvr diff --git a/app/shared_thread_data.py b/app/shared_thread_data.py new file mode 100644 index 0000000..40885e8 --- /dev/null +++ b/app/shared_thread_data.py @@ -0,0 +1,14 @@ +import threading + +class SharedThreadData: + def __init__(self, cfg): + self.transcript = "" + self.preview = "" + + self.stream = None + self.collector = None + + self.word_lock = threading.Lock() + self.exit_event = threading.Event() + self.cfg = cfg + diff --git a/app/steamvr.py b/app/steamvr.py new file mode 100644 index 0000000..64f34f5 --- /dev/null +++ b/app/steamvr.py @@ -0,0 +1,87 @@ +#!/usr/bin/env python3 + +import openvr as vr +import sys +import time + +EVENT_NONE = 0 +EVENT_RISING_EDGE = 1 +EVENT_FALLING_EDGE = 2 + +class InputEvent: + def __init__(self, + opcode: int): + self.opcode = opcode + +# Checks if the given button on the given controller is pressed. +def pollButtonPress( + hand: str = "right", + button: str = "b", + shared_data = None # SharedThreadData object + ) -> int: + hands = {} + hands["left"] = vr.TrackedControllerRole_LeftHand + hands["right"] = vr.TrackedControllerRole_RightHand + + buttons = {} + buttons["a"] = vr.k_EButton_IndexController_A + buttons["b"] = vr.k_EButton_IndexController_B + buttons["thumbstick"] = vr.k_EButton_Axis0 + + system = None + first = True + while not shared_data.exit_event.is_set() and not system: + try: + system = vr.init(vr.VRApplication_Background) + except Exception as e: + if first: + print(f"Failed to start steamVR input thread: {repr(e)}", file=sys.stderr) + first = False + time.sleep(1) + last_packet = 0 + event_high = False + + while not shared_data.exit_event.is_set(): + time.sleep(0.01) + + lh_idx = system.getTrackedDeviceIndexForControllerRole(hands[hand]) + #print("left hand device idx: {}".format(lh_idx)) + + got_state, state = system.getControllerState(lh_idx) + if not got_state: + continue + + if state.unPacketNum == last_packet: + continue + + # Clicking joysticks and moving joysticks fire the same events. To + # differentiate movement from clicking, we create a dead zone: if the event + # fires while the stick isn't moved far from center, we assume it's a + # click, not movement. + dead_zone_radius = 0.7 + + button_mask = (1 << buttons[button]) + ret = EVENT_NONE + if (state.ulButtonPressed & button_mask) != 0 and\ + (state.rAxis[0].x**2 + state.rAxis[0].y**2 < dead_zone_radius**2): + #print("button pressed: %016x" % state.ulButtonPressed) + #for i in range(0, 5): + # print("axis {} x: {} y: {}".format(i, state.rAxis[i].x, state.rAxis[i].y)) + if not event_high: + yield InputEvent(EVENT_RISING_EDGE) + event_high = True + elif event_high: + event_high = False + yield InputEvent(EVENT_FALLING_EDGE) + +if __name__ == "__main__": + gen = pollButtonPress() + while True: + time.sleep(0.1) + + event = pollButtonPress(session_state) + if event == EVENT_RISING_EDGE: + print("rising edge") + elif event == EVENT_FALLING_EDGE: + print("falling edge") + diff --git a/app/stt.py b/app/stt.py new file mode 100644 index 0000000..18f0f60 --- /dev/null +++ b/app/stt.py @@ -0,0 +1,915 @@ +from datetime import datetime +from faster_whisper import WhisperModel +import json +import langcodes +import numpy as np +import os +import noisereduce as nr +try: + from profanity_filter import ProfanityFilter + PROFANITY_FILTER_AVAILABLE = True +except ImportError: + PROFANITY_FILTER_AVAILABLE = False + print("Warning: profanity_filter module not available", file=sys.stderr) +import pyaudio +from pydub import AudioSegment +from shared_thread_data import SharedThreadData +from silero_vad import load_silero_vad, get_speech_timestamps +import sys +import time +import typing +import wave + +APP_ROOT = os.path.dirname(os.path.abspath(__file__)) +PROJECT_ROOT = os.path.dirname(APP_ROOT) + +class AudioStream(): + FORMAT = pyaudio.paInt16 + # Size of each frame (audio sample), in bytes. If you change FORMAT, make + # sure this stays up to date! + FRAME_SZ = 2 + # Frames per second. + FPS = 16000 + CHANNELS = 1 + def __init__(self): + pass + + def getSamples(self) -> bytes: + raise NotImplementedError("getSamples is not implemented!") + +class MicStream(AudioStream): + CHUNK_SZ = 1024 + + def __init__(self, cfg: typing.Dict): + self.p = pyaudio.PyAudio() + self.stream = None + self.sample_rate = None + # Each time pyaudio gives us audio data, it's in the form of a chunk of + # samples. We keep these in a list to keep the audio callback as light + # as possible. Whenever downstream layers want data, we collapse the + # list into a single array of data (a bytes object). + self.chunks = [] + # If set, incoming frames are simply discarded. + self.paused = False + + which_mic = cfg["microphone"] + + if cfg["enable_debug_mode"]: + print(f"Finding mic {which_mic}", file=sys.stderr) + self.dumpMicDevices() + + got_match = False + device_index = -1 + if which_mic == "index": + target_str = "Digital Audio Interface" + elif which_mic == "focusrite": + target_str = "Focusrite" + elif which_mic == "motu": + target_str = "In 1-2 (MOTU M Series)" + elif which_mic == "beyond": + target_str = "Microphone (Beyond)" + else: + if cfg["enable_debug_mode"]: + print(f"Mic {which_mic} requested, treating it as a numerical " + + "device ID", file=sys.stderr) + device_index = int(which_mic) + got_match = True + if not got_match: + info = self.p.get_host_api_info_by_index(0) + numdevices = info.get('deviceCount') + for i in range(0, numdevices): + if (self.p.get_device_info_by_host_api_device_index(0, i).get('maxInputChannels')) > 0: + device_name = self.p.get_device_info_by_host_api_device_index(0, i).get('name') + if target_str in device_name: + print(f"Got matching mic: {device_name}", + file=sys.stderr) + device_index = i + got_match = True + break + if not got_match: + raise KeyError(f"Mic {which_mic} not found") + + info = self.p.get_device_info_by_host_api_device_index(0, device_index) + if cfg["enable_debug_mode"]: + print(f"Found mic {which_mic}: {info['name']}", file=sys.stderr) + self.sample_rate = int(info['defaultSampleRate']) + if cfg["enable_debug_mode"]: + print(f"Mic sample rate: {self.sample_rate}", file=sys.stderr) + + self.stream = self.p.open( + rate=self.sample_rate, + channels=AudioStream.CHANNELS, + format=AudioStream.FORMAT, + input=True, + frames_per_buffer=MicStream.CHUNK_SZ, + input_device_index=device_index, + stream_callback=self.onAudioFramesAvailable) + + self.stream.start_stream() + + AudioStream.__init__(self) + + def pause(self, state: bool = True): + self.paused = state + + def dumpMicDevices(self): + info = self.p.get_host_api_info_by_index(0) + numdevices = info.get('deviceCount') + + for i in range(0, numdevices): + if (self.p.get_device_info_by_host_api_device_index(0, i).get('maxInputChannels')) > 0: + device_name = self.p.get_device_info_by_host_api_device_index(0, i).get('name') + print("Input Device id ", i, " - ", device_name) + + def onAudioFramesAvailable(self, + frames, + frame_count, + time_info, + status_flags): + if self.paused: + # Don't literally pause, just start returning silence. This allows + # the `min_segment_age_s` check to work while paused. + n_frames = int(frame_count * AudioStream.FPS / + float(self.sample_rate)) + self.chunks.append(np.zeros(n_frames, + dtype=np.int16).tobytes()) + return (frames, pyaudio.paContinue) + + decimated = b'' + # In pyaudio, a `frame` is a single sample of audio data. + frame_len = AudioStream.FRAME_SZ + next_frame = 0.0 + # The mic probably has a higher sample rate than Whisper wants, so + # decrease the sample rate by dropping samples. Note that this + # algorithm only works if the mic's rate is higher than whisper's + # expected rate. + keep_every = float(self.sample_rate) / AudioStream.FPS + for i in range(frame_count): + if i >= next_frame: + decimated += frames[i*frame_len:(i+1)*frame_len] + next_frame += keep_every + self.chunks.append(decimated) + + return (frames, pyaudio.paContinue) + + # Get audio data and the corresponding timestamp. + def getSamples(self) -> bytes: + chunks = self.chunks + self.chunks = [] + result = b''.join(chunks) + return result + +class AudioCollector: + def __init__(self, stream: AudioStream): + self.stream = stream + self.frames = b'' + # Note: by design, this is the only spot where we anchor our timestamps + # against the real world. This is done to make it possible to profile + # test cases which read from disk (at much faster than real speed) in + # the same way that we profile real-time data. + self.wall_ts = time.time() + + def getAudio(self) -> bytes: + frames = self.stream.getSamples() + if frames: + self.frames += frames + return self.frames + + def dropAudioPrefix(self, dur_s: float) -> bytes: + n_bytes = int(dur_s * AudioStream.FPS) * self.stream.FRAME_SZ + n_bytes = min(n_bytes, len(self.frames)) + cut_portion = self.frames[:n_bytes] + self.frames = self.frames[n_bytes:] + self.wall_ts += float(n_bytes / self.stream.FRAME_SZ) / self.stream.FPS + return cut_portion + + def dropAudioPrefixByFrames(self, dur_frames: int) -> bytes: + n_bytes = dur_frames * self.stream.FRAME_SZ + n_bytes = min(n_bytes, len(self.frames)) + cut_portion = self.frames[:n_bytes] + self.frames = self.frames[n_bytes:] + self.wall_ts += float(n_bytes / self.stream.FRAME_SZ) / self.stream.FPS + return cut_portion + + def keepLast(self, dur_s: float) -> bytes: + drop_len = max(0, self.duration() - dur_s) + return self.dropAudioPrefix(drop_len) + + def dropAudio(self): + self.wall_ts += self.duration() + cut_portion = self.frames + self.frames = b'' + return cut_portion + + def duration(self): + return len(self.frames) / (AudioStream.FPS * self.stream.FRAME_SZ) + + def begin(self): + return self.wall_ts + + def now(self): + return self.begin() + self.duration() + +class AudioCollectorFilter: + def __init__(self, parent: AudioCollector): + self.parent = parent + + def getAudio(self) -> bytes: + return self.parent.getAudio() + def dropAudioPrefix(self, dur_s: float): + return self.parent.dropAudioPrefix(dur_s) + def dropAudioPrefixByFrames(self, dur_frames: int): + return self.parent.dropAudioPrefixByFrames(dur_frames) + def keepLast(self, dur_s): + return self.parent.keepLast(dur_s) + def dropAudio(self): + return self.parent.dropAudio() + def duration(self): + return self.parent.duration() + def begin(self): + return self.parent.begin() + def now(self): + return self.parent.now() + +# Audio collector that enforces a minimum length on its audio data. +class LengthEnforcingAudioCollector(AudioCollectorFilter): + def __init__(self, parent: AudioCollector, min_duration_s: float): + AudioCollectorFilter.__init__(self, parent) + self.min_duration_s = min_duration_s + + def getAudio(self) -> bytes: + audio = self.parent.getAudio() + min_duration_frames = int(self.min_duration_s * AudioStream.FPS) + pad_len_frames = max(0, min_duration_frames - int(len(audio) / + AudioStream.FRAME_SZ)) + pad = np.zeros(pad_len_frames, dtype=np.int16).tobytes() + return pad + audio + +class NormalizingAudioCollector(AudioCollectorFilter): + def __init__(self, parent: AudioCollector): + AudioCollectorFilter.__init__(self, parent) + + def getAudio(self) -> bytes: + audio = self.parent.getAudio() + + audio = AudioSegment(audio, sample_width=AudioStream.FRAME_SZ, + frame_rate=AudioStream.FPS, channels=AudioStream.CHANNELS) + audio = audio.normalize() + + frames = np.array(audio.get_array_of_samples()) + frames = np.int16(frames).tobytes() + + return frames + +class BoostingAudioCollector(AudioCollectorFilter): + def __init__(self, parent: AudioCollector, + target_dBFS: float, + max_gain_dB: float, + cfg: typing.Dict): + AudioCollectorFilter.__init__(self, parent) + self.target_dBFS = target_dBFS + self.max_gain_dB = max_gain_dB + self.cfg = cfg + + def getAudio(self) -> bytes: + audio = self.parent.getAudio() + + audio = AudioSegment(audio, sample_width=AudioStream.FRAME_SZ, + frame_rate=AudioStream.FPS, channels=AudioStream.CHANNELS) + gain = min(self.target_dBFS - audio.dBFS, self.max_gain_dB) + if self.cfg["enable_debug_mode"]: + print(f"Boosting audio by {gain} dB (from {audio.dBFS} to {audio.dBFS + gain})", flush=True) + audio = audio.apply_gain(gain) + + frames = np.array(audio.get_array_of_samples()) + frames = np.int16(frames).tobytes() + + return frames + +class CompressingAudioCollector(AudioCollectorFilter): + def __init__(self, parent: AudioCollector): + AudioCollectorFilter.__init__(self, parent) + + def getAudio(self) -> bytes: + audio = self.parent.getAudio() + + audio = AudioSegment(audio, sample_width=AudioStream.FRAME_SZ, + frame_rate=AudioStream.FPS, channels=AudioStream.CHANNELS) + # subtle compression has a slight positive effect on my benchmark + audio = audio.compress_dynamic_range(threshold=-10, ratio=2.0) + + frames = np.array(audio.get_array_of_samples()) + frames = np.int16(frames).tobytes() + + return frames + +class NoiseReducingAudioCollector(AudioCollectorFilter): + def __init__(self, parent: AudioCollector, cfg: typing.Dict): + AudioCollectorFilter.__init__(self, parent) + self.cfg = cfg + + def getAudio(self) -> bytes: + audio = self.parent.getAudio() + audio_array = np.frombuffer(audio, dtype=np.int16).astype(np.float32) + + reduced_audio = nr.reduce_noise( + y=audio_array, + sr=AudioStream.FPS, + ) + + # Convert back to int16 + reduced_audio = np.clip(reduced_audio, -32768, 32767) + frames = np.int16(reduced_audio).tobytes() + + return frames + +class AudioSegmenter: + def __init__(self, + min_silence_ms=250, + max_speech_s=5, + min_speech_duration_ms=100): + self.min_silence_ms = min_silence_ms + self.max_speech_s = max_speech_s + self.min_speech_duration_ms = min_speech_duration_ms + + # Load Silero VAD model + self.model = load_silero_vad() + + self.vad_threshold = 0.3 + self.min_silence_duration_ms = min_silence_ms + self.max_speech_duration_s = max_speech_s + self.min_speech_duration_ms = min_speech_duration_ms + + def segmentAudio(self, audio: bytes): + # Convert audio bytes to numpy array expected by silero-vad + audio_array = np.frombuffer(audio, + dtype=np.int16).flatten().astype(np.float32) / 32768.0 + + # Get speech timestamps using silero-vad + # Note: silero-vad expects sample rate of 16000 Hz which matches AudioStream.FPS + speech_timestamps = get_speech_timestamps( + audio_array, + self.model, + sampling_rate=AudioStream.FPS, + threshold=self.vad_threshold, + min_silence_duration_ms=self.min_silence_duration_ms, + max_speech_duration_s=self.max_speech_duration_s, + min_speech_duration_ms=self.min_speech_duration_ms, + return_seconds=False # We want frame indices, not seconds + ) + + return speech_timestamps + + # Returns the stable cutoff (if any) and whether there are any segments. + def getStableCutoff(self, audio: bytes) -> typing.Tuple[int, bool]: + min_delta_frames = int((self.min_silence_duration_ms * + AudioStream.FPS) / 1000.0) + cutoff = None + + last_end = None + segments = self.segmentAudio(audio) + + for i in range(len(segments)): + s = segments[i] + #print(f"s: {s}") + #print(f"last_end: {last_end}") + + if last_end: + delta_frames = s['start'] - last_end + #print(f"delta frames: {delta_frames}") + if delta_frames > min_delta_frames: + cutoff = s['start'] + else: + last_end = s['end'] + + if i == len(segments) - 1: + now = int(len(audio) / AudioStream.FRAME_SZ) + #print(f"now: {now}") + #print(f"min d: {min_delta_frames}") + delta_frames = now - s['end'] + if delta_frames > min_delta_frames: + cutoff = now - int(min_delta_frames / 2) + + return (cutoff, len(segments) > 0) + +# A segment of transcribed audio. `start_ts` and `end_ts` are floating point +# number of seconds since the beginning of audio data. +class Segment: + def __init__(self, + transcript: str, + start_ts: float, + end_ts: float, + wall_ts: float, + avg_logprob: float, + no_speech_prob: float, + compression_ratio: float): + self.transcript = transcript + # start_ts, end_ts are timestamps in seconds relative to `wall_ts`. + self.start_ts = start_ts + self.end_ts = end_ts + # wall_ts is the time.time() at which the oldest audio sample leading + # to this transcript was collected. + self.wall_ts = wall_ts + self.avg_logprob = avg_logprob + self.no_speech_prob = no_speech_prob + self.compression_ratio = compression_ratio + + def __str__(self): + ts = f"(ts: {self.start_ts}-{self.end_ts}) " + + wall_ts_start = datetime.utcfromtimestamp(self.start_ts + self.wall_ts).strftime('%H:%M:%S') + wall_ts_end = datetime.utcfromtimestamp(self.end_ts + self.wall_ts).strftime('%H:%M:%S') + wall_ts = f"(wall ts: {wall_ts_start}-{wall_ts_end}) " + + no_speech = f"(no_speech: {self.no_speech_prob}) " + avg_logprob = f"(avg_logprob: {self.avg_logprob}) " + return f"{self.transcript} " + ts + wall_ts + no_speech + avg_logprob + +def join_segments(a, b): + if len(a) > 0 and a[-1] != ' ': + return a + ' ' + b + else: + return a + b + +class Whisper: + def __init__(self, + collector: AudioCollector, + cfg: typing.Dict): + self.collector = collector + self.model = None + self.cfg = cfg + + model_str = cfg["model"] + model_root = os.path.join(PROJECT_ROOT, "Models", + os.path.normpath(model_str)) + if cfg["enable_debug_mode"]: + print(f"Model {cfg['model']} will be saved to {model_root}", + file=sys.stderr) + + model_device = "cuda" + compute_type = cfg["compute_type"] + if cfg["use_cpu"]: + model_device = "cpu" + compute_type = "int8" + + already_downloaded = os.path.exists(model_root) + + if not already_downloaded: + print(f"Model {model_str} not already downloaded, downloading now...", flush=True) + + self.model = WhisperModel(model_str, + device = model_device, + device_index = cfg["gpu_idx"], + compute_type = compute_type, + download_root = model_root, + local_files_only = already_downloaded) + + self.context_window_chars = 200 # Keep last 200 chars of context + self.recent_context = "" # Store recent committed text + + def update_context(self, committed_text: str): + """Update the context with recently committed text.""" + self.recent_context = join_segments(self.recent_context, committed_text).strip() + # Drop half of the context window. + if len(self.recent_context) > self.context_window_chars: + words = self.recent_context.split() + words = words[len(words)//2:] + self.recent_context = ' '.join(words) + + def transcribe(self, frames: bytes = None) -> typing.List[Segment]: + if frames is None: + frames = self.collector.getAudio() + + # Convert audio to float32 + audio = np.frombuffer(frames, + dtype=np.int16).flatten().astype(np.float32) / 32768.0 + + # Build context-aware prompt + prompt = self._build_prompt() + + if self.cfg["enable_debug_mode"]: + print(f"Prompt: {prompt}", flush=True) + + t0 = time.time() + segments, info = self.model.transcribe( + audio, + language = langcodes.find(self.cfg["language"]).language, + vad_filter = True, + temperature=0.0, + without_timestamps = False, + initial_prompt=prompt, + beam_size=self.cfg.get("beam_size", 5), + best_of=self.cfg.get("best_of", 5), + condition_on_previous_text=True + ) + res = [] + for s in segments: + # Manual touchup. I see a decent number of hallucinations sneaking + # in with high `no_speech_prob` and modest `avg_logprob`. + if s.no_speech_prob > 0.6 and s.avg_logprob < -0.5: + if self.cfg["enable_debug_mode"]: + print(f"Drop probable hallucination (case 1) " + + f"(text='{s.text}', " + + f"no_speech_prob={s.no_speech_prob}, " + + f"avg_logprob={s.avg_logprob})", file=sys.stderr) + continue + # Another touchup targeted at the vexatious "thanks for watching!" + # hallucination. This triggers a lot when listening to + # instrumental/electronic music. + if s.no_speech_prob > 0.15 and s.avg_logprob < -0.7: + if self.cfg["enable_debug_mode"]: + print(f"Drop probable hallucination (case 2) " + + f"(text='{s.text}', " + + f"no_speech_prob={s.no_speech_prob}, " + + f"avg_logprob={s.avg_logprob})", file=sys.stderr) + continue + if s.avg_logprob < -0.75: + if self.cfg["enable_debug_mode"]: + print(f"Drop probable hallucination (case 3) " + + f"(text='{s.text}', " + + f"no_speech_prob={s.no_speech_prob}, " + + f"avg_logprob={s.avg_logprob})", file=sys.stderr) + continue + if self.cfg["enable_debug_mode"]: + print(f"s get: {s}") + if s.avg_logprob < -1.0: + continue + if s.compression_ratio > 2.4: + continue + res.append(Segment(s.text, s.start, s.end, + self.collector.begin(), + s.avg_logprob, s.no_speech_prob, s.compression_ratio)) + t1 = time.time() + if self.cfg["enable_debug_mode"]: + print(f"Transcription latency (s): {t1 - t0}") + return res + + def _build_prompt(self) -> str: + """Build a context-aware prompt for Whisper.""" + user_prompt = self.cfg["user_prompt"] + context_prompt = "" + if self.recent_context and len(self.recent_context) > 0: + context_prompt = f"Here is the context so far: {self.recent_context}" + + prompts = [user_prompt, context_prompt] + prompts = [p for p in prompts if p and len(p) > 0] + return " ".join(prompts) + +class TranscriptCommit: + def __init__(self, + delta: str, + preview: str, + latency_s: float = None, + thresh_at_commit: int = None, + audio: bytes = None, + duration_s: float = None, + start_ts: float = None): + self.delta = delta + self.preview = preview + self.latency_s = latency_s + self.thresh_at_commit = thresh_at_commit + self.audio = audio + # Time at which the commit is generated + self.ts = time.time() + # Time corresponding to the start of the segment + self.start_ts = start_ts + # The duration of the audio segment, in seconds. + self.duration_s = duration_s + + +def saveAudio(audio: bytes, path: str, cfg: typing.Dict): + with wave.open(path, 'wb') as wf: + if cfg["enable_debug_mode"]: + print(f"Saving audio to {path}", file=sys.stderr) + wf.setnchannels(AudioStream.CHANNELS) + wf.setsampwidth(AudioStream.FRAME_SZ) + wf.setframerate(AudioStream.FPS) + wf.writeframes(audio) + + +class SegmentLogger: + def __init__(self, cfg: typing.Dict): + self.cfg = cfg + self.enabled = cfg.get("enable_segment_logging", False) + self.session_data = [] + self.log_file = None + + if self.enabled: + log_dir = os.path.join(PROJECT_ROOT, "logs") + if not os.path.exists(log_dir): + os.makedirs(log_dir) + + # Create file + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + self.log_file = os.path.join(log_dir, f"session_debug_{timestamp}.json") + print(f"Segment logging enabled. Logging to: {self.log_file}", file=sys.stderr) + + def log_segment(self, segment: Segment, commit_type: str = "commit"): + if not self.enabled: + return + + segment_data = { + "timestamp": datetime.now().isoformat(), + "type": commit_type, + "text": segment.transcript, + "start_ts": segment.start_ts, + "end_ts": segment.end_ts, + "wall_ts": segment.wall_ts, + "avg_logprob": segment.avg_logprob, + "no_speech_prob": segment.no_speech_prob, + "compression_ratio": segment.compression_ratio, + "duration": segment.end_ts - segment.start_ts + } + + self.session_data.append(segment_data) + + # Write to file incrementally + try: + with open(self.log_file, 'w') as f: + json.dump({ + "session_start": self.session_data[0]["timestamp"] if self.session_data else None, + "segments": self.session_data + }, f, indent=2) + except Exception as e: + print(f"Error writing segment log: {e}", file=sys.stderr) + + def close(self): + if self.enabled and self.session_data: + print(f"Session complete. Logged {len(self.session_data)} segments to {self.log_file}", file=sys.stderr) + + +class VadCommitter: + def __init__(self, + cfg: typing.Dict, + collector: AudioCollector, + whisper: Whisper, + segmenter: AudioSegmenter, + segment_logger: SegmentLogger = None): + self.cfg = cfg + self.collector = collector + self.whisper = whisper + self.segmenter = segmenter + self.segment_logger = segment_logger + + def getDelta(self) -> TranscriptCommit: + audio = self.collector.getAudio() + stable_cutoff, has_audio = self.segmenter.getStableCutoff(audio) + + delta = "" + commit_audio = None + latency_s = None + duration_s = self.collector.duration() + start_ts = self.collector.begin() + + if has_audio and stable_cutoff: + latency_s = self.collector.now() - self.collector.begin() + duration_s = stable_cutoff / AudioStream.FPS + start_ts = self.collector.begin() + + # Get the filtered audio first, then extract the portion we need + filtered_audio = self.collector.getAudio() + commit_audio = filtered_audio[:stable_cutoff * AudioStream.FRAME_SZ] + + # Now drop the prefix from the collector + self.collector.dropAudioPrefixByFrames(stable_cutoff) + + segments = self.whisper.transcribe(commit_audio) + delta = ''.join(s.transcript for s in segments) + + # Update whisper's context with the committed text + if delta.strip(): + self.whisper.update_context(delta.strip()) + + if self.segment_logger: + for s in segments: + self.segment_logger.log_segment(s, "commit") + + audio = self.collector.getAudio() + if self.cfg["enable_debug_mode"]: + for s in segments: + print(f"commit segment: {s}", file=sys.stderr) + if len(delta) > 0: + print(f"delta get: {delta}", file=sys.stderr) + + if self.cfg["save_audio"] and len(delta) > 0: + ts = datetime.fromtimestamp(self.collector.now() - latency_s) + filename = str(ts.strftime('%Y_%m_%d__%H-%M-%S')) + delta.strip() + ".wav" + audio_dir = os.path.join(PROJECT_ROOT, "audio") + if not os.path.exists(audio_dir): + os.makedirs(audio_dir) + saveAudio(commit_audio, os.path.join(audio_dir, filename), self.cfg) + + preview = "" + if self.cfg["enable_previews"] and has_audio: + segments = self.whisper.transcribe(audio) + preview = "".join(s.transcript for s in segments) + + if self.segment_logger: + for s in segments: + self.segment_logger.log_segment(s, "preview") + + if not has_audio: + self.collector.keepLast(1.0) + + return TranscriptCommit( + delta.strip(), + preview.strip(), + latency_s, + audio=audio, + duration_s=duration_s, + start_ts=start_ts) + + +class StreamingPlugin: + def __init__(self): + pass + + def transform(self, commit: TranscriptCommit) -> TranscriptCommit: + return commit + + def stop(self): + pass + + +class LowercasePlugin(StreamingPlugin): + def __init__(self, cfg): + self.cfg = cfg + + def transform(self, commit: TranscriptCommit) -> TranscriptCommit: + if self.cfg["enable_lowercase_filter"]: + commit.delta = commit.delta.lower() + commit.preview = commit.preview.lower() + return commit + + +class UppercasePlugin(StreamingPlugin): + def __init__(self, cfg): + self.cfg = cfg + + def transform(self, commit: TranscriptCommit) -> TranscriptCommit: + if self.cfg["enable_uppercase_filter"]: + commit.delta = commit.delta.upper() + commit.preview = commit.preview.upper() + return commit + + +class ProfanityPlugin(StreamingPlugin): + def __init__(self, cfg): + self.cfg = cfg + self.filter = None + if PROFANITY_FILTER_AVAILABLE and cfg["enable_profanity_filter"]: + en_profanity_path = os.path.join(PROJECT_ROOT, "Third_Party/Profanity/en") + try: + self.filter = ProfanityFilter(en_profanity_path) + self.filter.load() + except Exception as e: + print(f"Warning: Could not load profanity filter: {e}", file=sys.stderr) + self.filter = None + + def transform(self, commit: TranscriptCommit) -> TranscriptCommit: + if self.cfg["enable_profanity_filter"] and self.filter: + commit.delta = self.filter.filter(commit.delta) + commit.preview = self.filter.filter(commit.preview) + return commit + + +class PresentationFilter: + def __init__(self): + pass + + def transform(self, transcript: str, preview: str) -> typing.Tuple[str, str]: + return transcript, preview + + def stop(self): + pass + + +class TrailingPeriodFilter(PresentationFilter): + def __init__(self, cfg): + self.cfg = cfg + + def transform(self, transcript: str, preview: str) -> typing.Tuple[str, str]: + if self.cfg["remove_trailing_period"]: + def _remove_trailing_period(s: str) -> str: + if len(s) > 0 and s[-1] == '.' and not s.endswith("..."): + s = s[0:len(s)-1] + return s + if len(preview) == 0: + transcript = _remove_trailing_period(transcript) + else: + preview = _remove_trailing_period(preview) + return transcript, preview + + +def transcriptionThread(shared_data: SharedThreadData): + last_stable_commit = None + + stream = MicStream(shared_data.cfg) + collector = AudioCollector(stream) + collector = CompressingAudioCollector(collector) + collector = BoostingAudioCollector(collector, -16.0, 24.0, + shared_data.cfg) + collector = NoiseReducingAudioCollector(collector, shared_data.cfg) + #collector = NormalizingAudioCollector(collector) + whisper = Whisper(collector, shared_data.cfg) + segmenter = AudioSegmenter(min_silence_ms=shared_data.cfg["min_silence_duration_ms"], + max_speech_s=shared_data.cfg["max_speech_duration_s"], + min_speech_duration_ms=shared_data.cfg["min_speech_duration_ms"]) + + segment_logger = SegmentLogger(shared_data.cfg) + committer = VadCommitter(shared_data.cfg, collector, whisper, segmenter, segment_logger) + + plugins = [] + # plugins.append(TranslationPlugin(shared_data.cfg)) # Not implemented yet + plugins.append(UppercasePlugin(shared_data.cfg)) + plugins.append(LowercasePlugin(shared_data.cfg)) + plugins.append(ProfanityPlugin(shared_data.cfg)) + # plugins.append(UwuPlugin(shared_data.cfg)) # Not implemented yet + # plugins.append(BrowserSource(shared_data.cfg)) # Not implemented yet + + filters = [] + filters.append(TrailingPeriodFilter(shared_data.cfg)) + + transcript = "" + preview = "" + + with shared_data.word_lock: + shared_data.stream = stream + shared_data.collector = collector + + print(f"Ready to go!", flush=True) + + while not shared_data.exit_event.is_set(): + time.sleep(shared_data.cfg["transcription_loop_delay_ms"] / 1000.0); + + op = None + + commit = committer.getDelta() + + with shared_data.word_lock: + for plugin in plugins: + commit = plugin.transform(commit) + + if len(commit.delta) > 0 or len(commit.preview) > 0: + # Avoid re-sending text after long pauses + if shared_data.cfg["reset_after_silence_s"] > 0: + silence_duration = 0 + if last_stable_commit: + last_commit_end_ts = \ + last_stable_commit.start_ts + \ + last_stable_commit.duration_s + silence_duration = commit.start_ts - last_commit_end_ts + if silence_duration > shared_data.cfg["reset_after_silence_s"]: + if shared_data.cfg["enable_debug_mode"]: + print(f"Resetting transcript after {silence_duration}-second " + "silence", file=sys.stderr) + shared_data.transcript = "" + shared_data.preview = "" + whisper.recent_context = "" # Reset context too + if commit.delta: + last_stable_commit = commit + + # Hard-cap displayed transcript length to prevent + # runaway memory use in UI. Keep the full transcript to avoid + # breaking OSC pager. + if len(shared_data.transcript) >= 1024: + shared_data.transcript = shared_data.transcript[-512:] + shared_data.transcript = \ + join_segments(shared_data.transcript, commit.delta) + shared_data.preview = commit.preview + + for filt in filters: + shared_data.transcript, shared_data.preview = \ + filt.transform(shared_data.transcript, + shared_data.preview) + + try: + print(f"Transcript: {shared_data.transcript}", flush=True) + except UnicodeEncodeError: + print("Failed to encode transcript - discarding delta", + file=sys.stderr) + continue + try: + print(f"Preview: {shared_data.preview}", flush=True) + except UnicodeEncodeError: + print("Failed to encode preview - discarding", file=sys.stderr) + + if shared_data.cfg["enable_debug_mode"]: + print(f"commit latency: {commit.latency_s}", file=sys.stderr) + print(f"commit thresh: {commit.thresh_at_commit}", + file=sys.stderr) + + if len(shared_data.transcript) > 0 and \ + (not shared_data.transcript.endswith(' ')) and \ + (not commit.delta.startswith(' ')): + commit.delta = ' ' + commit.delta + if len(commit.delta) > 0 and \ + (not commit.delta.endswith(' ')) and \ + (not commit.preview.startswith(' ')): + commit.preview = ' ' + commit.preview + for plugin in plugins: + plugin.stop() + for filt in filters: + filt.stop() + segment_logger.close() + |
