summaryrefslogtreecommitdiffstats
path: root/app
diff options
context:
space:
mode:
Diffstat (limited to 'app')
-rw-r--r--app/app_config.py39
-rw-r--r--app/hi.py580
-rw-r--r--app/keybind_event_machine.py21
-rw-r--r--app/list_microphones.py24
-rw-r--r--app/profanity_filter.py43
-rw-r--r--app/requirements.txt12
-rw-r--r--app/shared_thread_data.py14
-rw-r--r--app/steamvr.py87
-rw-r--r--app/stt.py915
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()
+