Parcourir la source

Optimized CPU performance

Efren Yevale Varela il y a 1 mois
Parent
commit
3285c4e9ce

+ 8 - 3
README.md

@@ -117,6 +117,9 @@ compute_type = "int8"                 # "int8" (CPU) or "float16" (CUDA, slower
 language     = "auto"                 # "auto" (detect), "en", "es", or None/""/"none"
 beam_size    = 1                      # 1 = fast greedy, >1 = slower beam search
 vad_filter   = false                  # Use Whisper's built-in VAD (we already have webrtcvad)
+cpu_threads  = 0                      # CPU inference threads (0 = CTranslate2's automatic choice).
+                                      # On many-core CPUs, set this to the physical core count
+                                      # (e.g. 16 on a Ryzen 9 8945HX) - measurably faster than auto.
 no_speech_prob_max = 0.6              # Drop segments Whisper scores as non-speech above this
 avg_logprob_min    = -1.0             # Drop segments with confidence below this
 
@@ -576,9 +579,11 @@ Then route audio into VirtualMicSink and try again.
 
 By default, the daemon uses `compute_type = "int8"` and `device = "cpu"`, which trades speed for accuracy. To speed up:
 
-1. **Reduce model size**: `--model base` or `--model tiny`
-2. **Enable GPU** (if you have CUDA): `--device cuda` with `compute_type float16`
-3. **Reduce beam_size**: Edit config to `beam_size = 1` (already default, very fast)
+1. **Pin `cpu_threads` to the physical core count**: `cpu_threads = 16` in `[whisper]`
+   (CTranslate2's automatic choice is often slower than an explicit count on many-core CPUs)
+2. **Reduce model size**: `--model base` or `--model tiny`
+3. **Enable GPU** (if you have CUDA): `--device cuda` with `compute_type float16`
+4. **Reduce beam_size**: Edit config to `beam_size = 1` (already default, very fast)
 
 Example:
 

+ 1 - 0
config.example.toml

@@ -25,6 +25,7 @@ compute_type       = "int8"
 language           = "auto"
 beam_size          = 1
 vad_filter         = false
+cpu_threads        = 0
 no_speech_prob_max = 0.6
 avg_logprob_min    = -1.0
 

+ 1 - 0
src/stt_daemon/cli.py

@@ -428,6 +428,7 @@ def _check(args) -> int:
             language=cfg.whisper.language,
             beam_size=cfg.whisper.beam_size,
             vad_filter=cfg.whisper.vad_filter,
+            cpu_threads=cfg.whisper.cpu_threads,
             no_speech_prob_max=cfg.whisper.no_speech_prob_max,
             avg_logprob_min=cfg.whisper.avg_logprob_min,
         )

+ 2 - 0
src/stt_daemon/config.py

@@ -56,6 +56,7 @@ class WhisperConfig:
     language: str | None = None
     beam_size: int = 1
     vad_filter: bool = False
+    cpu_threads: int = 0
     no_speech_prob_max: float = 0.6
     avg_logprob_min: float = -1.0
 
@@ -223,6 +224,7 @@ def load_config(path: str | None = None) -> Config:
                 language=language,
                 beam_size=whisper_data.get("beam_size", cfg.whisper.beam_size),
                 vad_filter=whisper_data.get("vad_filter", cfg.whisper.vad_filter),
+                cpu_threads=whisper_data.get("cpu_threads", cfg.whisper.cpu_threads),
                 no_speech_prob_max=whisper_data.get(
                     "no_speech_prob_max", cfg.whisper.no_speech_prob_max
                 ),

+ 1 - 0
src/stt_daemon/daemon.py

@@ -105,6 +105,7 @@ class Daemon:
             language=self.cfg.whisper.language,
             beam_size=self.cfg.whisper.beam_size,
             vad_filter=self.cfg.whisper.vad_filter,
+            cpu_threads=self.cfg.whisper.cpu_threads,
             no_speech_prob_max=self.cfg.whisper.no_speech_prob_max,
             avg_logprob_min=self.cfg.whisper.avg_logprob_min,
         )

+ 12 - 1
src/stt_daemon/transcriber.py

@@ -37,6 +37,7 @@ class Transcriber:
         language: str | None,
         beam_size: int = 1,
         vad_filter: bool = False,
+        cpu_threads: int = 0,
         no_speech_prob_max: float = 0.6,
         avg_logprob_min: float = -1.0,
     ):
@@ -49,6 +50,9 @@ class Transcriber:
             language: Language code (e.g., "en", "es") or None for auto-detect.
             beam_size: Beam search width (1 = greedy, fastest).
             vad_filter: Enable Whisper's built-in VAD (we already run webrtcvad).
+            cpu_threads: CPU threads for inference (0 = CTranslate2's automatic
+                choice). Setting this to the physical core count is typically
+                faster than the automatic choice on many-core CPUs.
             no_speech_prob_max: Drop segments Whisper scores as non-speech above this.
             avg_logprob_min: Drop segments with average log-probability below this.
         """
@@ -57,6 +61,7 @@ class Transcriber:
         self.compute_type = compute_type
         self.beam_size = beam_size
         self.vad_filter = vad_filter
+        self.cpu_threads = cpu_threads
         self.no_speech_prob_max = no_speech_prob_max
         self.avg_logprob_min = avg_logprob_min
 
@@ -76,7 +81,12 @@ class Transcriber:
                 self.device = "cpu"
                 self.compute_type = "int8"
 
-        self.model = WhisperModel(model, device=self.device, compute_type=self.compute_type)
+        self.model = WhisperModel(
+            model,
+            device=self.device,
+            compute_type=self.compute_type,
+            cpu_threads=self.cpu_threads,
+        )
 
     def transcribe(self, audio: np.ndarray) -> TranscriptResult:
         """Transcribe audio to text.
@@ -93,6 +103,7 @@ class Transcriber:
             beam_size=self.beam_size,
             vad_filter=self.vad_filter,
             condition_on_previous_text=False,
+            without_timestamps=True,
         )
 
         # Join the segments Whisper is confident about

+ 28 - 0
tests/test_config.py

@@ -263,3 +263,31 @@ require_viewer = false
 
     assert "assitant" in caplog.text
     assert cfg.assistant.require_viewer is True
+
+
+def test_cpu_threads_default_is_auto():
+    """cpu_threads defaults to 0 (CTranslate2's automatic choice)."""
+    with tempfile.TemporaryDirectory() as tmpdir:
+        import os
+        old_cwd = os.getcwd()
+        try:
+            os.chdir(tmpdir)
+            cfg = load_config()
+        finally:
+            os.chdir(old_cwd)
+
+    assert cfg.whisper.cpu_threads == 0
+
+
+def test_cpu_threads_config_field():
+    config_toml = """
+[whisper]
+cpu_threads = 16
+"""
+    with tempfile.TemporaryDirectory() as tmpdir:
+        config_path = Path(tmpdir) / "config.toml"
+        config_path.write_text(config_toml)
+
+        cfg = load_config(str(config_path))
+
+    assert cfg.whisper.cpu_threads == 16

+ 74 - 1
tests/test_transcriber.py

@@ -1,6 +1,10 @@
 """Tests for the Whisper segment confidence gate."""
 
-from stt_daemon.transcriber import accept_segment
+import numpy as np
+import pytest
+
+from stt_daemon import transcriber as transcriber_mod
+from stt_daemon.transcriber import Transcriber, accept_segment
 
 
 class _Seg:
@@ -31,3 +35,72 @@ def test_missing_attributes_default_to_accept():
         text = "x"
 
     assert accept_segment(Bare(), 0.6, -1.0)
+
+
+class _Info:
+    language = "en"
+    language_probability = 1.0
+
+
+class _FakeWhisperModel:
+    """Records constructor/transcribe kwargs without loading a real model."""
+
+    last = None
+
+    def __init__(self, model, device, compute_type, cpu_threads=0):
+        self.init_kwargs = {
+            "model": model,
+            "device": device,
+            "compute_type": compute_type,
+            "cpu_threads": cpu_threads,
+        }
+        self.transcribe_kwargs = None
+        _FakeWhisperModel.last = self
+
+    def transcribe(self, audio, **kwargs):
+        self.transcribe_kwargs = kwargs
+        return [], _Info()
+
+
+@pytest.fixture
+def fake_model(monkeypatch):
+    monkeypatch.setattr(transcriber_mod, "WhisperModel", _FakeWhisperModel)
+    return _FakeWhisperModel
+
+
+def _audio():
+    return np.zeros(16000, dtype=np.float32)
+
+
+def test_cpu_threads_forwarded_to_model(fake_model):
+    t = Transcriber(
+        "small", device="cpu", compute_type="int8", language=None, cpu_threads=16
+    )
+    assert t.cpu_threads == 16
+    assert fake_model.last.init_kwargs["cpu_threads"] == 16
+
+
+def test_cpu_threads_default_is_auto(fake_model):
+    Transcriber("small", device="cpu", compute_type="int8", language=None)
+    assert fake_model.last.init_kwargs["cpu_threads"] == 0
+
+
+def test_transcribe_disables_timestamps(fake_model):
+    t = Transcriber("small", device="cpu", compute_type="int8", language=None)
+    t.transcribe(_audio())
+    kwargs = fake_model.last.transcribe_kwargs
+    assert kwargs["without_timestamps"] is True
+    assert kwargs["condition_on_previous_text"] is False
+
+
+def test_auto_language_stays_none_for_bilingual(fake_model):
+    """language=auto must reach the engine as None so es/en both keep working."""
+    t = Transcriber("small", device="cpu", compute_type="int8", language=None)
+    t.transcribe(_audio())
+    assert fake_model.last.transcribe_kwargs["language"] is None
+
+
+def test_explicit_language_is_preserved(fake_model):
+    t = Transcriber("small", device="cpu", compute_type="int8", language="es")
+    t.transcribe(_audio())
+    assert fake_model.last.transcribe_kwargs["language"] == "es"