More polish

- Filters actually get applied now, huge accuracy boost
- Use silero-vad python library instead of rolling our own
- Expose prompt parameter
- Auto setup venv on launch
- Clean up python output
- Auto acquire all dependencies on launch
- Add icon

4 files changed, 109 insertions, 347 deletions

diff --git a/app/hi.py b/app/hi.py
index 0d80b9d..e6877ff 100644
--- a/app/hi.py
+++ b/app/hi.py
@@ -330,10 +330,11 @@ 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()
+    if False:
+        osc_thread = threading.Thread(
+                target=osc_thread,
+                args=(shared_data,))
+        osc_thread.start()
 
     transcribe_thread = threading.Thread(
             target=stt.transcriptionThread,
@@ -382,6 +383,7 @@ if __name__ == "__main__":
                 local_word = shared_data.word
         print(local_word + "_")
     shared_data.exit_event.set()
-    osc_thread.join()
+    if False:
+        osc_thread.join()
     transcribe_thread.join()
 
diff --git a/app/requirements.txt b/app/requirements.txt
index 07f94cd..f8b7069 100644
--- a/app/requirements.txt
+++ b/app/requirements.txt
@@ -5,4 +5,4 @@ pyaudio
 pydub
 python-osc
 sentencepiece
-wave
+silero-vad
diff --git a/app/stt.py b/app/stt.py
index c157f6d..7d76333 100644
--- a/app/stt.py
+++ b/app/stt.py
@@ -6,10 +6,10 @@ import os
 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 vad
 import wave
 
 
@@ -33,7 +33,7 @@ class AudioStream():
 class MicStream(AudioStream):
     CHUNK_SZ = 1024
 
-    def __init__(self, which_mic: str):
+    def __init__(self, cfg: typing.Dict):
         self.p = pyaudio.PyAudio()
         self.stream = None
         self.sample_rate = None
@@ -45,8 +45,11 @@ class MicStream(AudioStream):
         # If set, incoming frames are simply discarded.
         self.paused = False
 
-        print(f"Finding mic {which_mic}", file=sys.stderr)
-        self.dumpMicDevices()
+        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
@@ -59,8 +62,9 @@ class MicStream(AudioStream):
         elif which_mic == "beyond":
             target_str = "Microphone (Beyond)"
         else:
-            print(f"Mic {which_mic} requested, treating it as a numerical " +
-                    "device ID", file=sys.stderr)
+            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:
@@ -79,9 +83,11 @@ class MicStream(AudioStream):
             raise KeyError(f"Mic {which_mic} not found")
 
         info = self.p.get_device_info_by_host_api_device_index(0, device_index)
-        print(f"Found mic {which_mic}: {info['name']}", file=sys.stderr)
+        if cfg["enable_debug_mode"]:
+            print(f"Found mic {which_mic}: {info['name']}", file=sys.stderr)
         self.sample_rate = int(info['defaultSampleRate'])
-        print(f"Mic sample rate: {self.sample_rate}", file=sys.stderr)
+        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,
@@ -289,19 +295,40 @@ class AudioSegmenter:
     def __init__(self,
             min_silence_ms=250,
             max_speech_s=5):
-        self.vad_options = vad.VadOptions(
-                min_silence_duration_ms=min_silence_ms,
-                max_speech_duration_s=max_speech_s)
-        pass
+        self.min_silence_ms = min_silence_ms
+        self.max_speech_s = max_speech_s
+        
+        # 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.speech_pad_ms = 300
 
     def segmentAudio(self, audio: bytes):
-        audio = np.frombuffer(audio,
+        # Convert audio bytes to numpy array expected by silero-vad
+        audio_array = np.frombuffer(audio,
                 dtype=np.int16).flatten().astype(np.float32) / 32768.0
-        return vad.get_speech_timestamps(audio, vad_options=self.vad_options)
+        
+        # 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,
+            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.vad_options.min_silence_duration_ms *
+        min_delta_frames = int((self.min_silence_duration_ms *
                 AudioStream.FPS) / 1000.0)
         cutoff = None
 
@@ -379,8 +406,9 @@ class Whisper:
         model_str = cfg["model"]
         model_root = os.path.join(parent_dir, "Models",
                 os.path.normpath(model_str))
-        print(f"Model {cfg['model']} will be saved to {model_root}",
-                file=sys.stderr)
+        if cfg["enable_debug_mode"]:
+            print(f"Model {cfg['model']} will be saved to {model_root}",
+                    file=sys.stderr)
 
         model_device = "cuda"
         if cfg["use_cpu"]:
@@ -395,21 +423,42 @@ class Whisper:
                 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 = (self.recent_context + " " + committed_text).strip()
+        # Keep only the last N characters to avoid prompt getting too long
+        if len(self.recent_context) > self.context_window_chars:
+            self.recent_context = self.recent_context[-self.context_window_chars:]
+
     def transcribe(self, frames: bytes = None) -> typing.List[Segment]:
         if frames is None:
             frames = self.collector.getAudio()
-        # Convert from signed 16-bit int [-32768, 32767] to signed 32-bit float on
-        # [-1, 1].
+        
+        # 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()
+        
         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)
+                without_timestamps = False,
+                initial_prompt=prompt,
+                beam_size=5,
+                best_of=5,
+                condition_on_previous_text=True,
+                compression_ratio_threshold=2.4,
+                log_prob_threshold=-1.0,
+                no_speech_threshold=0.6
+        )
         res = []
         for s in segments:
             # Manual touchup. I see a decent number of hallucinations sneaking
@@ -445,6 +494,17 @@ class Whisper:
             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,
@@ -502,10 +562,21 @@ class VadCommitter:
             latency_s = self.collector.now() - self.collector.begin()
             duration_s = stable_cutoff / AudioStream.FPS
             start_ts = self.collector.begin()
-            commit_audio = self.collector.dropAudioPrefixByFrames(stable_cutoff)
+            
+            # 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())
+            
             audio = self.collector.getAudio()
             if self.cfg["enable_debug_mode"]:
                 for s in segments:
@@ -540,11 +611,11 @@ class VadCommitter:
 def transcriptionThread(shared_data: SharedThreadData):
     last_stable_commit = None
 
-    stream = MicStream(shared_data.cfg["microphone"])
+    stream = MicStream(shared_data.cfg)
     collector = AudioCollector(stream)
     collector = CompressingAudioCollector(collector)
+    collector = BoostingAudioCollector(collector, -12.0, shared_data.cfg)
     collector = NormalizingAudioCollector(collector)
-    collector = BoostingAudioCollector(collector, 0.0, shared_data.cfg)
     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"])
@@ -553,6 +624,8 @@ def transcriptionThread(shared_data: SharedThreadData):
     transcript = ""
     preview = ""
 
+    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);
 
@@ -561,8 +634,7 @@ def transcriptionThread(shared_data: SharedThreadData):
         commit = committer.getDelta()
 
         if len(commit.delta) > 0 or len(commit.preview) > 0:
-            # Avoid re-sending text after long pauses. User controls the length
-            # of the pause in the UI.
+            # Avoid re-sending text after long pauses
             if shared_data.cfg["reset_after_silence_s"] > 0:
                 silence_duration = 0
                 if last_stable_commit:
@@ -571,10 +643,12 @@ def transcriptionThread(shared_data: SharedThreadData):
                             last_stable_commit.duration_s
                     silence_duration = commit.start_ts - last_commit_end_ts
                 if silence_duration > shared_data.cfg["reset_after_silence_s"]:
-                    print(f"Resetting transcript after {silence_duration}-second "
-                            "silence", file=sys.stderr)
+                    if shared_data.cfg["enable_debug_mode"]:
+                        print(f"Resetting transcript after {silence_duration}-second "
+                                "silence", file=sys.stderr)
                     transcript = ""
                     preview = ""
+                    whisper.recent_context = ""  # Reset context too
                 if commit.delta:
                     last_stable_commit = commit
 
diff --git a/app/vad.py b/app/vad.py
deleted file mode 100644
index 1dea765..0000000
--- a/app/vad.py
+++ /dev/null
@@ -1,314 +0,0 @@
-# MIT License
-# 
-# Copyright (c) 2023 Guillaume Klein
-# 
-# Permission is hereby granted, free of charge, to any person obtaining a copy
-# of this software and associated documentation files (the "Software"), to deal
-# in the Software without restriction, including without limitation the rights
-# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-# copies of the Software, and to permit persons to whom the Software is
-# furnished to do so, subject to the following conditions:
-# 
-# The above copyright notice and this permission notice shall be included in all
-# copies or substantial portions of the Software.
-# 
-# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-# SOFTWARE.
-
-import bisect
-import functools
-import os
-import warnings
-
-from typing import List, NamedTuple, Optional
-
-import numpy as np
-
-
-# The code below is adapted from https://github.com/snakers4/silero-vad.
-class VadOptions(NamedTuple):
-    """VAD options.
-
-    Attributes:
-      threshold: Speech threshold. Silero VAD outputs speech probabilities for each audio chunk,
-        probabilities ABOVE this value are considered as SPEECH. It is better to tune this
-        parameter for each dataset separately, but "lazy" 0.5 is pretty good for most datasets.
-      min_speech_duration_ms: Final speech chunks shorter min_speech_duration_ms are thrown out.
-      max_speech_duration_s: Maximum duration of speech chunks in seconds. Chunks longer
-        than max_speech_duration_s will be split at the timestamp of the last silence that
-        lasts more than 100ms (if any), to prevent aggressive cutting. Otherwise, they will be
-        split aggressively just before max_speech_duration_s.
-      min_silence_duration_ms: In the end of each speech chunk wait for min_silence_duration_ms
-        before separating it
-      window_size_samples: Audio chunks of window_size_samples size are fed to the silero VAD model.
-        WARNING! Silero VAD models were trained using 512, 1024, 1536 samples for 16000 sample rate.
-        Values other than these may affect model performance!!
-      speech_pad_ms: Final speech chunks are padded by speech_pad_ms each side
-    """
-
-    threshold: float = 0.5
-    min_speech_duration_ms: int = 250
-    max_speech_duration_s: float = float("inf")
-    min_silence_duration_ms: int = 2000
-    window_size_samples: int = 1024
-    speech_pad_ms: int = 400
-
-
-def get_speech_timestamps(
-    audio: np.ndarray,
-    vad_options: Optional[VadOptions] = None,
-    **kwargs,
-) -> List[dict]:
-    """This method is used for splitting long audios into speech chunks using silero VAD.
-
-    Args:
-      audio: One dimensional float array.
-      vad_options: Options for VAD processing.
-      kwargs: VAD options passed as keyword arguments for backward compatibility.
-
-    Returns:
-      List of dicts containing begin and end samples of each speech chunk.
-    """
-    if vad_options is None:
-        vad_options = VadOptions(**kwargs)
-
-    threshold = vad_options.threshold
-    min_speech_duration_ms = vad_options.min_speech_duration_ms
-    max_speech_duration_s = vad_options.max_speech_duration_s
-    min_silence_duration_ms = vad_options.min_silence_duration_ms
-    window_size_samples = vad_options.window_size_samples
-    speech_pad_ms = vad_options.speech_pad_ms
-
-    if window_size_samples not in [512, 1024, 1536]:
-        warnings.warn(
-            "Unusual window_size_samples! Supported window_size_samples:\n"
-            " - [512, 1024, 1536] for 16000 sampling_rate"
-        )
-
-    sampling_rate = 16000
-    min_speech_samples = sampling_rate * min_speech_duration_ms / 1000
-    speech_pad_samples = sampling_rate * speech_pad_ms / 1000
-    max_speech_samples = (
-        sampling_rate * max_speech_duration_s
-        - window_size_samples
-        - 2 * speech_pad_samples
-    )
-    min_silence_samples = sampling_rate * min_silence_duration_ms / 1000
-    min_silence_samples_at_max_speech = sampling_rate * 98 / 1000
-
-    audio_length_samples = len(audio)
-
-    model = get_vad_model()
-    state = model.get_initial_state(batch_size=1)
-
-    speech_probs = []
-    for current_start_sample in range(0, audio_length_samples, window_size_samples):
-        chunk = audio[current_start_sample : current_start_sample + window_size_samples]
-        if len(chunk) < window_size_samples:
-            chunk = np.pad(chunk, (0, int(window_size_samples - len(chunk))))
-        speech_prob, state = model(chunk, state, sampling_rate)
-        speech_probs.append(speech_prob)
-
-    triggered = False
-    speeches = []
-    current_speech = {}
-    neg_threshold = threshold - 0.15
-
-    # to save potential segment end (and tolerate some silence)
-    temp_end = 0
-    # to save potential segment limits in case of maximum segment size reached
-    prev_end = next_start = 0
-
-    for i, speech_prob in enumerate(speech_probs):
-        if (speech_prob >= threshold) and temp_end:
-            temp_end = 0
-            if next_start < prev_end:
-                next_start = window_size_samples * i
-
-        if (speech_prob >= threshold) and not triggered:
-            triggered = True
-            current_speech["start"] = window_size_samples * i
-            continue
-
-        if (
-            triggered
-            and (window_size_samples * i) - current_speech["start"] > max_speech_samples
-        ):
-            if prev_end:
-                current_speech["end"] = prev_end
-                speeches.append(current_speech)
-                current_speech = {}
-                # previously reached silence (< neg_thres) and is still not speech (< thres)
-                if next_start < prev_end:
-                    triggered = False
-                else:
-                    current_speech["start"] = next_start
-                prev_end = next_start = temp_end = 0
-            else:
-                current_speech["end"] = window_size_samples * i
-                speeches.append(current_speech)
-                current_speech = {}
-                prev_end = next_start = temp_end = 0
-                triggered = False
-                continue
-
-        if (speech_prob < neg_threshold) and triggered:
-            if not temp_end:
-                temp_end = window_size_samples * i
-            # condition to avoid cutting in very short silence
-            if (window_size_samples * i) - temp_end > min_silence_samples_at_max_speech:
-                prev_end = temp_end
-            if (window_size_samples * i) - temp_end < min_silence_samples:
-                continue
-            else:
-                current_speech["end"] = temp_end
-                if (
-                    current_speech["end"] - current_speech["start"]
-                ) > min_speech_samples:
-                    speeches.append(current_speech)
-                current_speech = {}
-                prev_end = next_start = temp_end = 0
-                triggered = False
-                continue
-
-    if (
-        current_speech
-        and (audio_length_samples - current_speech["start"]) > min_speech_samples
-    ):
-        current_speech["end"] = audio_length_samples
-        speeches.append(current_speech)
-
-    for i, speech in enumerate(speeches):
-        if i == 0:
-            speech["start"] = int(max(0, speech["start"] - speech_pad_samples))
-        if i != len(speeches) - 1:
-            silence_duration = speeches[i + 1]["start"] - speech["end"]
-            if silence_duration < 2 * speech_pad_samples:
-                speech["end"] += int(silence_duration // 2)
-                speeches[i + 1]["start"] = int(
-                    max(0, speeches[i + 1]["start"] - silence_duration // 2)
-                )
-            else:
-                speech["end"] = int(
-                    min(audio_length_samples, speech["end"] + speech_pad_samples)
-                )
-                speeches[i + 1]["start"] = int(
-                    max(0, speeches[i + 1]["start"] - speech_pad_samples)
-                )
-        else:
-            speech["end"] = int(
-                min(audio_length_samples, speech["end"] + speech_pad_samples)
-            )
-
-    return speeches
-
-
-def collect_chunks(audio: np.ndarray, chunks: List[dict]) -> np.ndarray:
-    """Collects and concatenates audio chunks."""
-    if not chunks:
-        return np.array([], dtype=np.float32)
-
-    return np.concatenate([audio[chunk["start"] : chunk["end"]] for chunk in chunks])
-
-
-class SpeechTimestampsMap:
-    """Helper class to restore original speech timestamps."""
-
-    def __init__(self, chunks: List[dict], sampling_rate: int, time_precision: int = 2):
-        self.sampling_rate = sampling_rate
-        self.time_precision = time_precision
-        self.chunk_end_sample = []
-        self.total_silence_before = []
-
-        previous_end = 0
-        silent_samples = 0
-
-        for chunk in chunks:
-            silent_samples += chunk["start"] - previous_end
-            previous_end = chunk["end"]
-
-            self.chunk_end_sample.append(chunk["end"] - silent_samples)
-            self.total_silence_before.append(silent_samples / sampling_rate)
-
-    def get_original_time(
-        self,
-        time: float,
-        chunk_index: Optional[int] = None,
-    ) -> float:
-        if chunk_index is None:
-            chunk_index = self.get_chunk_index(time)
-
-        total_silence_before = self.total_silence_before[chunk_index]
-        return round(total_silence_before + time, self.time_precision)
-
-    def get_chunk_index(self, time: float) -> int:
-        sample = int(time * self.sampling_rate)
-        return min(
-            bisect.bisect(self.chunk_end_sample, sample),
-            len(self.chunk_end_sample) - 1,
-        )
-
-
-@functools.lru_cache
-def get_vad_model():
-    """Returns the VAD model instance."""
-    abspath = os.path.abspath(__file__)
-    my_dir = os.path.dirname(abspath)
-    parent_dir = os.path.dirname(my_dir)
-    path = os.path.join(parent_dir, "Models", "silero_vad.onnx")
-    return SileroVADModel(path)
-
-
-class SileroVADModel:
-    def __init__(self, path):
-        try:
-            import onnxruntime
-        except ImportError as e:
-            raise RuntimeError(
-                "Applying the VAD filter requires the onnxruntime package"
-            ) from e
-
-        opts = onnxruntime.SessionOptions()
-        opts.inter_op_num_threads = 1
-        opts.intra_op_num_threads = 1
-        opts.log_severity_level = 4
-
-        self.session = onnxruntime.InferenceSession(
-            path,
-            providers=["CPUExecutionProvider"],
-            sess_options=opts,
-        )
-
-    def get_initial_state(self, batch_size: int):
-        h = np.zeros((2, batch_size, 64), dtype=np.float32)
-        c = np.zeros((2, batch_size, 64), dtype=np.float32)
-        return h, c
-
-    def __call__(self, x, state, sr: int):
-        if len(x.shape) == 1:
-            x = np.expand_dims(x, 0)
-        if len(x.shape) > 2:
-            raise ValueError(
-                f"Too many dimensions for input audio chunk {len(x.shape)}"
-            )
-        if sr / x.shape[1] > 31.25:
-            raise ValueError("Input audio chunk is too short")
-
-        h, c = state
-
-        ort_inputs = {
-            "input": x,
-            "h": h,
-            "c": c,
-            "sr": np.array(sr, dtype="int64"),
-        }
-
-        out, h, c = self.session.run(None, ort_inputs)
-        state = (h, c)
-
-        return out, state