yum-slop/TaSTT

Free self-hosted STT for VRChat.

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

yumDrop turbo; use old logic when no_speech ts availablebce0853

master
33.1 KiB904 linesraw
1from datetime import datetime
2from faster_whisper import WhisperModel
3import json
4import langcodes
5from logger import log, log_err
6import noisereduce as nr
7import numpy as np
8import os
9from profanity_filter import ProfanityFilter
10import pyaudio
11from pydub import AudioSegment
12from shared_thread_data import SharedThreadData
13from silero_vad import load_silero_vad, get_speech_timestamps
14import sys
15import time
16import typing
17import wave
18from hallucination_filter import HallucinationFilter
19
20APP_ROOT = os.path.dirname(os.path.abspath(__file__))
21PROJECT_ROOT = os.path.dirname(APP_ROOT)
22
23class AudioStream():
24    FORMAT = pyaudio.paInt16
25    # Size of each frame (audio sample), in bytes. If you change FORMAT, make
26    # sure this stays up to date!
27    FRAME_SZ = 2
28    # Frames per second.
29    FPS = 16000
30    CHANNELS = 1
31    def __init__(self):
32        pass
33
34    def getSamples(self) -> bytes:
35        raise NotImplementedError("getSamples is not implemented!")
36
37class MicStream(AudioStream):
38    CHUNK_SZ = 1024
39
40    def __init__(self, cfg: typing.Dict):
41        self.p = pyaudio.PyAudio()
42        self.stream = None
43        self.sample_rate = None
44        # Each time pyaudio gives us audio data, it's in the form of a chunk of
45        # samples. We keep these in a list to keep the audio callback as light
46        # as possible. Whenever downstream layers want data, we collapse the
47        # list into a single array of data (a bytes object).
48        self.chunks = []
49        # If set, incoming frames are simply discarded.
50        self.paused = False
51
52        which_mic = cfg["microphone"]
53
54        if cfg["enable_debug_mode"]:
55            log(f"Finding mic {which_mic}")
56            self.dumpMicDevices()
57
58        got_match = False
59        device_index = -1
60        if which_mic == "index":
61            target_str = "Digital Audio Interface"
62        elif which_mic == "focusrite":
63            target_str = "Focusrite"
64        elif which_mic == "motu":
65            target_str = "In 1-2 (MOTU M Series)"
66        elif which_mic == "beyond":
67            target_str = "Microphone (Beyond)"
68        else:
69            if cfg["enable_debug_mode"]:
70                log(f"Mic {which_mic} requested, treating it as a numerical " +
71                        "device ID")
72            device_index = int(which_mic)
73            got_match = True
74        if not got_match:
75            info = self.p.get_host_api_info_by_index(0)
76            numdevices = info.get('deviceCount')
77            for i in range(0, numdevices):
78                if (self.p.get_device_info_by_host_api_device_index(0, i).get('maxInputChannels')) > 0:
79                    device_name = self.p.get_device_info_by_host_api_device_index(0, i).get('name')
80                    if target_str in device_name:
81                        log(f"Got matching mic: {device_name}")
82                        device_index = i
83                        got_match = True
84                        break
85        if not got_match:
86            raise KeyError(f"Mic {which_mic} not found")
87
88        info = self.p.get_device_info_by_host_api_device_index(0, device_index)
89        if cfg["enable_debug_mode"]:
90            log(f"Found mic {which_mic}: {info['name']}")
91        self.sample_rate = int(info['defaultSampleRate'])
92        if cfg["enable_debug_mode"]:
93            log(f"Mic sample rate: {self.sample_rate}")
94
95        self.stream = self.p.open(
96                rate=self.sample_rate,
97                channels=AudioStream.CHANNELS,
98                format=AudioStream.FORMAT,
99                input=True,
100                frames_per_buffer=MicStream.CHUNK_SZ,
101                input_device_index=device_index,
102                stream_callback=self.onAudioFramesAvailable)
103
104        self.stream.start_stream()
105
106        AudioStream.__init__(self)
107
108    def pause(self, state: bool = True):
109        self.paused = state
110
111    def dumpMicDevices(self):
112        info = self.p.get_host_api_info_by_index(0)
113        numdevices = info.get('deviceCount')
114
115        for i in range(0, numdevices):
116            if (self.p.get_device_info_by_host_api_device_index(0, i).get('maxInputChannels')) > 0:
117                device_name = self.p.get_device_info_by_host_api_device_index(0, i).get('name')
118                log("Input Device id ", i, " - ", device_name)
119
120    def onAudioFramesAvailable(self,
121            frames,
122            frame_count,
123            time_info,
124            status_flags):
125        if self.paused:
126            # Don't literally pause, just start returning silence. This allows
127            # the `min_segment_age_s` check to work while paused.
128            n_frames = int(frame_count * AudioStream.FPS /
129                    float(self.sample_rate))
130            self.chunks.append(np.zeros(n_frames,
131                dtype=np.int16).tobytes())
132            return (frames, pyaudio.paContinue)
133
134        decimated = b''
135        # In pyaudio, a `frame` is a single sample of audio data.
136        frame_len = AudioStream.FRAME_SZ
137        next_frame = 0.0
138        # The mic probably has a higher sample rate than Whisper wants, so
139        # decrease the sample rate by dropping samples. Note that this
140        # algorithm only works if the mic's rate is higher than whisper's
141        # expected rate.
142        keep_every = float(self.sample_rate) / AudioStream.FPS
143        for i in range(frame_count):
144            if i >= next_frame:
145                decimated += frames[i*frame_len:(i+1)*frame_len]
146                next_frame += keep_every
147        self.chunks.append(decimated)
148
149        return (frames, pyaudio.paContinue)
150
151    # Get audio data and the corresponding timestamp.
152    def getSamples(self) -> bytes:
153        chunks = self.chunks
154        self.chunks = []
155        result = b''.join(chunks)
156        return result
157
158class AudioCollector:
159    def __init__(self, stream: AudioStream):
160        self.stream = stream
161        self.frames = b''
162        # Note: by design, this is the only spot where we anchor our timestamps
163        # against the real world. This is done to make it possible to profile
164        # test cases which read from disk (at much faster than real speed) in
165        # the same way that we profile real-time data.
166        self.wall_ts = time.time()
167
168    def getAudio(self) -> bytes:
169        frames = self.stream.getSamples()
170        if frames:
171            self.frames += frames
172        return self.frames
173
174    def dropAudioPrefix(self, dur_s: float) -> bytes:
175        n_bytes = int(dur_s * AudioStream.FPS) * self.stream.FRAME_SZ
176        n_bytes = min(n_bytes, len(self.frames))
177        cut_portion = self.frames[:n_bytes]
178        self.frames = self.frames[n_bytes:]
179        self.wall_ts += float(n_bytes / self.stream.FRAME_SZ) / self.stream.FPS
180        return cut_portion
181
182    def dropAudioPrefixByFrames(self, dur_frames: int) -> bytes:
183        n_bytes = dur_frames * self.stream.FRAME_SZ
184        n_bytes = min(n_bytes, len(self.frames))
185        cut_portion = self.frames[:n_bytes]
186        self.frames = self.frames[n_bytes:]
187        self.wall_ts += float(n_bytes / self.stream.FRAME_SZ) / self.stream.FPS
188        return cut_portion
189
190    def keepLast(self, dur_s: float) -> bytes:
191        drop_len = max(0, self.duration() - dur_s)
192        return self.dropAudioPrefix(drop_len)
193
194    def dropAudio(self):
195        self.wall_ts += self.duration()
196        cut_portion = self.frames
197        self.frames = b''
198        return cut_portion
199
200    def duration(self):
201        return len(self.frames) / (AudioStream.FPS * self.stream.FRAME_SZ)
202
203    def begin(self):
204        return self.wall_ts
205
206    def now(self):
207        return self.begin() + self.duration()
208
209class AudioCollectorFilter:
210    def __init__(self, parent: AudioCollector):
211        self.parent = parent
212
213    def getAudio(self) -> bytes:
214        return self.parent.getAudio()
215    def dropAudioPrefix(self, dur_s: float):
216        return self.parent.dropAudioPrefix(dur_s)
217    def dropAudioPrefixByFrames(self, dur_frames: int):
218        return self.parent.dropAudioPrefixByFrames(dur_frames)
219    def keepLast(self, dur_s):
220        return self.parent.keepLast(dur_s)
221    def dropAudio(self):
222        return self.parent.dropAudio()
223    def duration(self):
224        return self.parent.duration()
225    def begin(self):
226        return self.parent.begin()
227    def now(self):
228        return self.parent.now()
229
230# Audio collector that enforces a minimum length on its audio data.
231class LengthEnforcingAudioCollector(AudioCollectorFilter):
232    def __init__(self, parent: AudioCollector, min_duration_s: float):
233        AudioCollectorFilter.__init__(self, parent)
234        self.min_duration_s = min_duration_s
235
236    def getAudio(self) -> bytes:
237        audio = self.parent.getAudio()
238        min_duration_frames = int(self.min_duration_s * AudioStream.FPS)
239        pad_len_frames = max(0, min_duration_frames - int(len(audio) /
240            AudioStream.FRAME_SZ))
241        pad = np.zeros(pad_len_frames, dtype=np.int16).tobytes()
242        return pad + audio
243
244class NormalizingAudioCollector(AudioCollectorFilter):
245    def __init__(self, parent: AudioCollector):
246        AudioCollectorFilter.__init__(self, parent)
247
248    def getAudio(self) -> bytes:
249        audio = self.parent.getAudio()
250
251        audio = AudioSegment(audio, sample_width=AudioStream.FRAME_SZ,
252                frame_rate=AudioStream.FPS, channels=AudioStream.CHANNELS)
253        audio = audio.normalize()
254
255        frames = np.array(audio.get_array_of_samples())
256        frames = np.int16(frames).tobytes()
257
258        return frames
259
260class BoostingAudioCollector(AudioCollectorFilter):
261    def __init__(self, parent: AudioCollector,
262                 target_dBFS: float,
263                 max_gain_dB: float,
264                 cfg: typing.Dict):
265        AudioCollectorFilter.__init__(self, parent)
266        self.target_dBFS = target_dBFS
267        self.max_gain_dB = max_gain_dB
268        self.cfg = cfg
269
270    def getAudio(self) -> bytes:
271        audio = self.parent.getAudio()
272
273        audio = AudioSegment(audio, sample_width=AudioStream.FRAME_SZ,
274                frame_rate=AudioStream.FPS, channels=AudioStream.CHANNELS)
275        gain = min(self.target_dBFS - audio.dBFS, self.max_gain_dB)
276        if self.cfg["enable_debug_mode"]:
277            log(f"Boosting audio by {gain} dB (from {audio.dBFS} to {audio.dBFS + gain})")
278        audio = audio.apply_gain(gain)
279
280        frames = np.array(audio.get_array_of_samples())
281        frames = np.int16(frames).tobytes()
282
283        return frames
284
285class CompressingAudioCollector(AudioCollectorFilter):
286    def __init__(self, parent: AudioCollector):
287        AudioCollectorFilter.__init__(self, parent)
288
289    def getAudio(self) -> bytes:
290        audio = self.parent.getAudio()
291
292        audio = AudioSegment(audio, sample_width=AudioStream.FRAME_SZ,
293                frame_rate=AudioStream.FPS, channels=AudioStream.CHANNELS)
294        # subtle compression has a slight positive effect on my benchmark
295        audio = audio.compress_dynamic_range(threshold=-10, ratio=2.0)
296
297        frames = np.array(audio.get_array_of_samples())
298        frames = np.int16(frames).tobytes()
299
300        return frames
301
302class NoiseReducingAudioCollector(AudioCollectorFilter):
303    def __init__(self, parent: AudioCollector, cfg: typing.Dict):
304        AudioCollectorFilter.__init__(self, parent)
305        self.cfg = cfg
306
307    def getAudio(self) -> bytes:
308        audio = self.parent.getAudio()
309        audio_array = np.frombuffer(audio, dtype=np.int16).astype(np.float32)
310
311        reduced_audio = nr.reduce_noise(
312            y=audio_array,
313            sr=AudioStream.FPS,
314        )
315
316        # Convert back to int16
317        reduced_audio = np.clip(reduced_audio, -32768, 32767)
318        frames = np.int16(reduced_audio).tobytes()
319
320        return frames
321
322class AudioSegmenter:
323    def __init__(self,
324            min_silence_ms=250,
325            max_speech_s=5,
326            min_speech_duration_ms=100):
327        self.min_silence_ms = min_silence_ms
328        self.max_speech_s = max_speech_s
329        self.min_speech_duration_ms = min_speech_duration_ms
330
331        # Load Silero VAD model
332        self.model = load_silero_vad()
333
334        self.min_silence_duration_ms = min_silence_ms
335        self.max_speech_duration_s = max_speech_s
336        self.min_speech_duration_ms = min_speech_duration_ms
337
338    def segmentAudio(self, audio: bytes):
339        # Convert audio bytes to numpy array expected by silero-vad
340        audio_array = np.frombuffer(audio,
341                dtype=np.int16).flatten().astype(np.float32) / 32768.0
342
343        # Get speech timestamps using silero-vad
344        # Note: silero-vad expects sample rate of 16000 Hz which matches AudioStream.FPS
345        speech_timestamps = get_speech_timestamps(
346            audio_array,
347            self.model,
348            sampling_rate=AudioStream.FPS,
349            min_silence_duration_ms=self.min_silence_duration_ms,
350            max_speech_duration_s=self.max_speech_duration_s,
351            min_speech_duration_ms=self.min_speech_duration_ms,
352            return_seconds=False  # We want frame indices, not seconds
353        )
354
355        return speech_timestamps
356
357    # Returns the stable cutoff (if any) and whether there are any segments.
358    def getStableCutoff(self, audio: bytes) -> typing.Tuple[int, bool]:
359        min_delta_frames = int((self.min_silence_duration_ms *
360                AudioStream.FPS) / 1000.0)
361        cutoff = None
362
363        last_end = None
364        segments = self.segmentAudio(audio)
365
366        for i in range(len(segments)):
367            s = segments[i]
368
369            if last_end:
370                delta_frames = s['start'] - last_end
371                if delta_frames > min_delta_frames:
372                    cutoff = s['start']
373            else:
374                last_end = s['end']
375
376            if i == len(segments) - 1:
377                now = int(len(audio) / AudioStream.FRAME_SZ)
378                delta_frames = now - s['end']
379                if delta_frames > min_delta_frames:
380                    cutoff = now - int(min_delta_frames / 2)
381
382        return (cutoff, len(segments) > 0)
383
384# A segment of transcribed audio. `start_ts` and `end_ts` are floating point
385# number of seconds since the beginning of audio data.
386class Segment:
387    def __init__(self,
388            transcript: str,
389            start_ts: float,
390            end_ts: float,
391            wall_ts: float,
392            avg_logprob: float,
393            no_speech_prob: float,
394            compression_ratio: float,
395            audio_len_s: float):
396        self.transcript = transcript
397        # start_ts, end_ts are timestamps in seconds relative to `wall_ts`.
398        self.start_ts = start_ts
399        self.end_ts = end_ts
400        # wall_ts is the time.time() at which the oldest audio sample leading
401        # to this transcript was collected.
402        self.wall_ts = wall_ts
403        self.avg_logprob = avg_logprob
404        self.no_speech_prob = no_speech_prob
405        self.compression_ratio = compression_ratio
406        self.audio_len_s = audio_len_s
407
408    def __str__(self):
409        ts = f"(ts: {self.start_ts}-{self.end_ts}) "
410
411        wall_ts_start = datetime.utcfromtimestamp(self.start_ts + self.wall_ts).strftime('%H:%M:%S')
412        wall_ts_end = datetime.utcfromtimestamp(self.end_ts + self.wall_ts).strftime('%H:%M:%S')
413        wall_ts = f"(wall ts: {wall_ts_start}-{wall_ts_end}) "
414
415        no_speech = f"(no_speech: {self.no_speech_prob}) "
416        avg_logprob = f"(avg_logprob: {self.avg_logprob}) "
417        max_len_s = f"(max_len_s: {self.audio_len_s}) "
418        return f"{self.transcript} " + ts + wall_ts + no_speech + avg_logprob + max_len_s
419
420def join_segments(a, b):
421    if len(a) > 0 and a[-1] != ' ':
422        return a + ' ' + b
423    else:
424        return a + b
425
426
427class SegmentLogger:
428    def __init__(self, cfg: typing.Dict):
429        self.cfg = cfg
430        self.enabled = cfg.get("enable_segment_logging", False)
431        self.session_data = []
432        self.log_file = None
433
434        if self.enabled:
435            log_dir = os.path.join(PROJECT_ROOT, "logs")
436            if not os.path.exists(log_dir):
437                os.makedirs(log_dir)
438
439            # Create file
440            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
441            self.log_file = os.path.join(log_dir, f"session_debug_{timestamp}.json")
442            log(f"Segment logging enabled. Logging to: {self.log_file}")
443
444    def log_segment(self, segment: Segment, commit_type: str = "commit"):
445        if not self.enabled:
446            return
447
448        segment_data = {
449            "timestamp": datetime.now().isoformat(),
450            "type": commit_type,
451            "text": segment.transcript,
452            "start_ts": segment.start_ts,
453            "end_ts": segment.end_ts,
454            "wall_ts": segment.wall_ts,
455            "avg_logprob": segment.avg_logprob,
456            "no_speech_prob": segment.no_speech_prob,
457            "compression_ratio": segment.compression_ratio,
458            "duration": segment.end_ts - segment.start_ts,
459            "duration_sanity": segment.audio_len_s
460        }
461
462        self.session_data.append(segment_data)
463
464        # Write to file incrementally
465        try:
466            with open(self.log_file, 'w') as f:
467                json.dump({
468                    "session_start": self.session_data[0]["timestamp"] if self.session_data else None,
469                    "segments": self.session_data
470                }, f, indent=2)
471        except Exception as e:
472            log_err(f"Error writing segment log: {e}")
473
474    def close(self):
475        if self.enabled and self.session_data:
476            log(f"Session complete. Logged {len(self.session_data)} " + \
477                    "segments to {self.log_file}")
478
479
480class Whisper:
481    def __init__(self,
482            collector: AudioCollector,
483            cfg: typing.Dict,
484            segment_logger: SegmentLogger = None):
485        self.collector = collector
486        self.model = None
487        self.cfg = cfg
488        self.hallucination_filter = HallucinationFilter(cfg)
489        self.segment_logger = segment_logger
490
491        model_str = cfg["model"]
492        model_root = os.path.join(PROJECT_ROOT, "Models",
493                os.path.normpath(model_str))
494        if cfg["enable_debug_mode"]:
495            log(f"Model {cfg['model']} will be saved to {model_root}")
496
497        model_device = "cuda"
498        compute_type = cfg["compute_type"]
499        if cfg["use_cpu"]:
500            model_device = "cpu"
501            compute_type = "int8"
502
503        already_downloaded = os.path.exists(model_root)
504
505        if not already_downloaded:
506            log(f"Model {model_str} not already downloaded, downloading now...")
507
508        self.model = WhisperModel(model_str,
509                device = model_device,
510                device_index = cfg["gpu_idx"],
511                compute_type = compute_type,
512                download_root = model_root,
513                local_files_only = already_downloaded)
514
515        self.context_window_chars = 200  # Keep last 200 chars of context
516        self.recent_context = ""  # Store recent committed text
517
518    def update_context(self, committed_text: str):
519        """Update the context with recently committed text."""
520        self.recent_context = join_segments(self.recent_context, committed_text).strip()
521        # Drop half of the context window.
522        if len(self.recent_context) > self.context_window_chars:
523            words = self.recent_context.split()
524            words = words[len(words)//2:]
525            self.recent_context = ' '.join(words)
526
527    def transcribe(self, frames: bytes = None) -> typing.List[Segment]:
528        if frames is None:
529            frames = self.collector.getAudio()
530
531        # Convert audio to float32
532        audio = np.frombuffer(frames,
533                dtype=np.int16).flatten().astype(np.float32) / 32768.0
534        audio_len_s = len(frames) / 16000.0
535
536        # Build context-aware prompt
537        prompt = self._build_prompt()
538
539        if self.cfg["enable_debug_mode"]:
540            log(f"Prompt: {prompt}")
541
542        t0 = time.time()
543        segments, info = self.model.transcribe(
544                audio,
545                language = langcodes.find(self.cfg["language"]).language,
546                vad_filter = False,
547                #temperature=0.0,
548                without_timestamps = False,
549                initial_prompt=prompt,
550                beam_size=self.cfg.get("beam_size", 5),
551                best_of=self.cfg.get("best_of", 5),
552                condition_on_previous_text=True
553        )
554        res = []
555        for s in segments:
556            # Log raw segment before filtering
557            if self.segment_logger:
558                # Create a temporary segment object for logging
559                raw_seg = Segment(s.text, s.start, s.end,
560                    self.collector.begin(),
561                    s.avg_logprob, s.no_speech_prob, s.compression_ratio,
562                    audio_len_s)
563                self.segment_logger.log_segment(raw_seg, "raw")
564            # Sometimes the model reports a bum duration, breaking our filters.
565            # Cap the segment length above by the length of the audio in.
566            duration_s = min(s.end - s.start, audio_len_s)
567
568            if self.cfg["enable_debug_mode"]:
569                log(f"s get: {s}")
570            if s.avg_logprob < -1.0:
571                continue
572            if s.compression_ratio > 2.4:
573                continue
574
575            # Create segment object
576            seg = Segment(s.text, s.start, s.end,
577                self.collector.begin(),
578                s.avg_logprob, s.no_speech_prob, s.compression_ratio,
579                audio_len_s)
580
581            # Apply hallucination filter. This is a statistical model trained
582            # on personal speech data.
583            if self.hallucination_filter.is_hallucination(seg):
584                if self.cfg["enable_debug_mode"]:
585                    log(f"Drop probable hallucination (case 4) " +
586                            f"(text='{s.text}', " +
587                            f"no_speech_prob={s.no_speech_prob}, " +
588                            f"avg_logprob={s.avg_logprob})")
589                continue
590
591            res.append(seg)
592        t1 = time.time()
593        if self.cfg["enable_debug_mode"]:
594            log(f"Transcription latency (s): {t1 - t0}")
595        return res
596
597    def _build_prompt(self) -> str:
598        """Build a context-aware prompt for Whisper."""
599        user_prompt = self.cfg["user_prompt"]
600        context_prompt = ""
601        if self.recent_context and len(self.recent_context) > 0:
602            context_prompt = f"Here is the context so far: {self.recent_context}"
603
604        prompts = [user_prompt, context_prompt]
605        prompts = [p for p in prompts if p and len(p) > 0]
606        return " ".join(prompts)
607
608class TranscriptCommit:
609    def __init__(self,
610            delta: str,
611            preview: str,
612            latency_s: float = None,
613            thresh_at_commit: int = None,
614            audio: bytes = None,
615            duration_s: float = None,
616            start_ts: float = None):
617        self.delta = delta
618        self.preview = preview
619        self.latency_s = latency_s
620        self.thresh_at_commit = thresh_at_commit
621        self.audio = audio
622        # Time at which the commit is generated
623        self.ts = time.time()
624        # Time corresponding to the start of the segment
625        self.start_ts = start_ts
626        # The duration of the audio segment, in seconds.
627        self.duration_s = duration_s
628
629
630def saveAudio(audio: bytes, path: str, cfg: typing.Dict):
631    with wave.open(path, 'wb') as wf:
632        if cfg["enable_debug_mode"]:
633            log(f"Saving audio to {path}")
634        wf.setnchannels(AudioStream.CHANNELS)
635        wf.setsampwidth(AudioStream.FRAME_SZ)
636        wf.setframerate(AudioStream.FPS)
637        wf.writeframes(audio)
638
639
640class VadCommitter:
641    def __init__(self,
642            cfg: typing.Dict,
643            collector: AudioCollector,
644            whisper: Whisper,
645            segmenter: AudioSegmenter,
646            segment_logger: SegmentLogger = None):
647        self.cfg = cfg
648        self.collector = collector
649        self.whisper = whisper
650        self.segmenter = segmenter
651        self.segment_logger = segment_logger
652
653    def getDelta(self) -> TranscriptCommit:
654        audio = self.collector.getAudio()
655        stable_cutoff, has_audio = self.segmenter.getStableCutoff(audio)
656
657        delta = ""
658        commit_audio = None
659        latency_s = None
660        duration_s = self.collector.duration()
661        start_ts = self.collector.begin()
662
663        if has_audio and stable_cutoff:
664            latency_s = self.collector.now() - self.collector.begin()
665            duration_s = stable_cutoff / AudioStream.FPS
666            start_ts = self.collector.begin()
667
668            # Get the filtered audio first, then extract the portion we need
669            filtered_audio = self.collector.getAudio()
670            commit_audio = filtered_audio[:stable_cutoff * AudioStream.FRAME_SZ]
671
672            # Now drop the prefix from the collector
673            self.collector.dropAudioPrefixByFrames(stable_cutoff)
674
675            segments = self.whisper.transcribe(commit_audio)
676            delta = ''.join(s.transcript for s in segments)
677
678            # Update whisper's context with the committed text
679            if delta.strip():
680                self.whisper.update_context(delta.strip())
681
682            audio = self.collector.getAudio()
683            if self.cfg["enable_debug_mode"]:
684                for s in segments:
685                    log(f"commit segment: {s}")
686                if len(delta) > 0:
687                    log(f"delta get: {delta}")
688
689            if self.cfg["save_audio"] and len(delta) > 0:
690                ts = datetime.fromtimestamp(self.collector.now() - latency_s)
691                sanitized_delta = ''.join(c for c in delta.strip() if c.isalnum() or c == ' ')
692                filename = str(ts.strftime('%Y_%m_%d__%H-%M-%S')) + sanitized_delta + ".wav"
693                audio_dir = os.path.join(PROJECT_ROOT, "audio")
694                if not os.path.exists(audio_dir):
695                    os.makedirs(audio_dir)
696                saveAudio(commit_audio, os.path.join(audio_dir, filename), self.cfg)
697
698        preview = ""
699        if self.cfg["enable_previews"] and has_audio:
700            segments = self.whisper.transcribe(audio)
701            preview = "".join(s.transcript for s in segments)
702
703        if not has_audio:
704            self.collector.keepLast(1.0)
705
706        return TranscriptCommit(
707                delta.strip(),
708                preview.strip(),
709                latency_s,
710                audio=audio,
711                duration_s=duration_s,
712                start_ts=start_ts)
713
714
715class StreamingPlugin:
716    def __init__(self):
717        pass
718
719    def transform(self, commit: TranscriptCommit) -> TranscriptCommit:
720        return commit
721
722    def stop(self):
723        pass
724
725
726class LowercasePlugin(StreamingPlugin):
727    def __init__(self, cfg):
728        self.cfg = cfg
729
730    def transform(self, commit: TranscriptCommit) -> TranscriptCommit:
731        if self.cfg["enable_lowercase_filter"]:
732            commit.delta = commit.delta.lower()
733            commit.preview = commit.preview.lower()
734        return commit
735
736
737class UppercasePlugin(StreamingPlugin):
738    def __init__(self, cfg):
739        self.cfg = cfg
740
741    def transform(self, commit: TranscriptCommit) -> TranscriptCommit:
742        if self.cfg["enable_uppercase_filter"]:
743            commit.delta = commit.delta.upper()
744            commit.preview = commit.preview.upper()
745        return commit
746
747
748class ProfanityPlugin(StreamingPlugin):
749    def __init__(self, cfg):
750        self.cfg = cfg
751        self.filter = None
752        if cfg["enable_profanity_filter"]:
753            en_profanity_path = os.path.join(PROJECT_ROOT, "Third_Party/Profanity/en")
754            try:
755                self.filter = ProfanityFilter(en_profanity_path)
756                self.filter.load()
757            except Exception as e:
758                log_err(f"Warning: Could not load profanity filter: {e}")
759                self.filter = None
760
761    def transform(self, commit: TranscriptCommit) -> TranscriptCommit:
762        if self.cfg["enable_profanity_filter"] and self.filter:
763            commit.delta = self.filter.filter(commit.delta)
764            commit.preview = self.filter.filter(commit.preview)
765        return commit
766
767
768class PresentationFilter:
769    def __init__(self):
770        pass
771
772    def transform(self, transcript: str, preview: str) -> typing.Tuple[str, str]:
773        return transcript, preview
774
775    def stop(self):
776        pass
777
778
779class TrailingPeriodFilter(PresentationFilter):
780    def __init__(self, cfg):
781        self.cfg = cfg
782
783    def transform(self, transcript: str, preview: str) -> typing.Tuple[str, str]:
784        if self.cfg["remove_trailing_period"]:
785            def _remove_trailing_period(s: str) -> str:
786                if len(s) > 0 and s[-1] == '.' and not s.endswith("..."):
787                    s = s[0:len(s)-1]
788                return s
789            if len(preview) == 0:
790                transcript = _remove_trailing_period(transcript)
791            else:
792                preview = _remove_trailing_period(preview)
793        return transcript, preview
794
795
796def transcriptionThread(shared_data: SharedThreadData):
797    last_stable_commit = None
798
799    stream = MicStream(shared_data.cfg)
800    collector = AudioCollector(stream)
801    collector = CompressingAudioCollector(collector)
802    collector = BoostingAudioCollector(collector, -16.0, 24.0,
803                                       shared_data.cfg)
804    collector = NoiseReducingAudioCollector(collector, shared_data.cfg)
805    #collector = NormalizingAudioCollector(collector)
806    segment_logger = SegmentLogger(shared_data.cfg)
807    whisper = Whisper(collector, shared_data.cfg, segment_logger)
808    segmenter = AudioSegmenter(min_silence_ms=shared_data.cfg["min_silence_duration_ms"],
809            max_speech_s=shared_data.cfg["max_speech_duration_s"],
810            min_speech_duration_ms=shared_data.cfg["min_speech_duration_ms"])
811    committer = VadCommitter(shared_data.cfg, collector, whisper, segmenter, segment_logger)
812
813    plugins = []
814    # plugins.append(TranslationPlugin(shared_data.cfg))  # Not implemented yet
815    plugins.append(UppercasePlugin(shared_data.cfg))
816    plugins.append(LowercasePlugin(shared_data.cfg))
817    plugins.append(ProfanityPlugin(shared_data.cfg))
818    # plugins.append(UwuPlugin(shared_data.cfg))  # Not implemented yet
819    # plugins.append(BrowserSource(shared_data.cfg))  # Not implemented yet
820
821    filters = []
822    filters.append(TrailingPeriodFilter(shared_data.cfg))
823
824    transcript = ""
825    preview = ""
826
827    with shared_data.word_lock:
828        shared_data.stream = stream
829        shared_data.collector = collector
830
831    log(f"Ready to go!")
832
833    while not shared_data.exit_event.is_set():
834        time.sleep(shared_data.cfg["transcription_loop_delay_ms"] / 1000.0);
835
836        op = None
837
838        commit = committer.getDelta()
839
840        with shared_data.word_lock:
841            for plugin in plugins:
842                commit = plugin.transform(commit)
843
844            if len(commit.delta) > 0 or len(commit.preview) > 0:
845                # Avoid re-sending text after long pauses
846                if shared_data.cfg["reset_after_silence_s"] > 0:
847                    silence_duration = 0
848                    if last_stable_commit:
849                        last_commit_end_ts = \
850                                last_stable_commit.start_ts + \
851                                last_stable_commit.duration_s
852                        silence_duration = commit.start_ts - last_commit_end_ts
853                    if silence_duration > shared_data.cfg["reset_after_silence_s"]:
854                        if shared_data.cfg["enable_debug_mode"]:
855                            log(f"Resetting transcript after {silence_duration}-second "
856                                    "silence")
857                        shared_data.transcript = ""
858                        shared_data.preview = ""
859                        whisper.recent_context = ""  # Reset context too
860                    if commit.delta:
861                        last_stable_commit = commit
862
863                # Hard-cap displayed transcript length to prevent
864                # runaway memory use in UI. Keep the full transcript to avoid
865                # breaking OSC pager.
866                if len(shared_data.transcript) >= 1024:
867                    shared_data.transcript = shared_data.transcript[-512:]
868                shared_data.transcript = \
869                        join_segments(shared_data.transcript, commit.delta)
870                shared_data.preview = commit.preview
871
872                for filt in filters:
873                    shared_data.transcript, shared_data.preview = \
874                            filt.transform(shared_data.transcript,
875                                           shared_data.preview)
876
877                try:
878                    log(f"Transcript: {shared_data.transcript}")
879                except UnicodeEncodeError:
880                    log_err("Failed to encode transcript - discarding delta")
881                    continue
882                try:
883                    log(f"Preview: {shared_data.preview}")
884                except UnicodeEncodeError:
885                    log_err("Failed to encode preview - discarding")
886
887                if shared_data.cfg["enable_debug_mode"]:
888                    log(f"commit latency: {commit.latency_s}")
889                    log(f"commit thresh: {commit.thresh_at_commit}")
890
891            if len(shared_data.transcript) > 0 and \
892                    (not shared_data.transcript.endswith(' ')) and \
893                    (not commit.delta.startswith(' ')):
894                commit.delta = ' ' + commit.delta
895            if len(commit.delta) > 0 and \
896                    (not commit.delta.endswith(' ')) and \
897                    (not commit.preview.startswith(' ')):
898                commit.preview = ' ' + commit.preview
899    for plugin in plugins:
900        plugin.stop()
901    for filt in filters:
902        filt.stop()
903    segment_logger.close()
904