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
branchesmasterall refs →
3.9 KiB103 linesraw
1import io
2import joblib
3from logger import log, log_err
4import numpy as np
5import pandas as pd
6from pathlib import Path
7import pronouncing
8import re
9import sys
10
11def count_syllables(word):
12    """Count syllables in a word using pronouncing library with regex fallback."""
13    phones = pronouncing.phones_for_word(word.lower())
14    if len(phones) == 0:
15        return 0
16    return pronouncing.syllable_count(phones[0])
17
18def text_syllable_count(text):
19    """Count total syllables in text."""
20    words = re.findall(r'\b\w+\b', text)
21    return sum(count_syllables(word) for word in words)
22
23class HallucinationFilter:
24    """Filter for detecting hallucinated segments in speech-to-text output."""
25    def __init__(self, cfg, model_path: Path = None):
26        """
27        Initialize the hallucination filter.
28        Args:
29            model_path: Optional path to the model file. If not provided,
30                       uses the default path.
31        """
32        self.cfg = cfg
33        self.model = None
34        self.threshold = None
35        self.features = None
36        # Get the project root directory
37        app_root = Path(__file__).resolve().parent
38        project_root = app_root.parent
39        model_path = project_root / "Models" / "thankyou_filter_gb.pkl"
40        # Try to load the model
41        log(f"Loading hallucination filter")
42        bundle = joblib.load(model_path)
43        self.model = bundle["model"]
44        self.threshold = bundle["threshold"]
45        self.features = bundle["features"]
46        log(f"Loaded hallucination filter model from {model_path}")
47    def is_hallucination(self, segment) -> bool:
48        """
49        Check if a segment is likely a hallucination.
50        Returns False if model is not available.
51        Args:
52            segment: A segment object with attributes avg_logprob, audio_len_s,
53                    no_speech_prob, compression_ratio, text, start, and end.
54        Returns:
55            bool: True if the segment is likely a hallucination, False otherwise.
56        """
57        s = segment  # Brevity
58
59        if s.no_speech_prob == 0:
60            # no_speech is not available. Use fancy classifier trained on my
61            # speech data.
62            text = s.transcript
63            duration = s.audio_len_s
64            raw_duration = s.end_ts - s.start_ts
65            n_syllables = text_syllable_count(text)
66            sps = n_syllables / duration
67            raw_sps = n_syllables / raw_duration
68            duration_ratio = raw_duration / duration
69            X = pd.DataFrame([[
70                s.avg_logprob,
71                s.no_speech_prob,
72                s.compression_ratio,
73                np.log1p(duration),
74                np.log1p(sps),
75                np.log1p(raw_duration),
76                np.log1p(raw_sps),
77                duration_ratio,
78                s.avg_logprob * duration
79            ]], columns=self.features)
80            # Get probability
81            prob = self.model.predict_proba(X)[0, 1]
82            return prob >= self.threshold
83
84        # If no_speech is set, use simpler filter.
85        if s.no_speech_prob > 0.6 and s.avg_logprob < -0.5:
86            if self.cfg["enable_debug_mode"]:
87                print(f"Drop probable hallucination (case 1) " +
88                        f"(text='{s.text}', " +
89                        f"no_speech_prob={s.no_speech_prob}, " +
90                        f"avg_logprob={s.avg_logprob})", file=sys.stderr)
91            return True
92        # Another touchup targeted at the vexatious "thanks for watching!"
93        # hallucination. This triggers a lot when listening to
94        # instrumental/electronic music.
95        if s.no_speech_prob > 0.15 and s.avg_logprob < -0.7:
96            if self.cfg["enable_debug_mode"]:
97                print(f"Drop probable hallucination (case 2) " +
98                        f"(text='{s.text}', " +
99                        f"no_speech_prob={s.no_speech_prob}, " +
100                        f"avg_logprob={s.avg_logprob})", file=sys.stderr)
101            return True
102        return False
103