Initial commit

transcriptapi/stt.py

@@ -0,0 +1,56 @@
+# Copyright (c) 2024 Julian Müller (ChaoticByte)
+
+from io import BytesIO as _BytesIO
+from pathlib import Path as _Path
+
+from faster_whisper import WhisperModel as _WhisperModel
+
+from .msg import ComponentLogger as _ComponentLogger
+
+
+class STT:
+
+    def __init__(self, model_path: _Path, n_threads: int = 4, use_int8: bool = True, logger: _ComponentLogger = _ComponentLogger("STT")):
+        assert isinstance(model_path, _Path)
+        assert type(n_threads) == int and n_threads >= 0
+        assert type(use_int8) == bool
+        assert isinstance(logger, _ComponentLogger)
+        self.logger = logger
+        self.model_path = model_path.expanduser().resolve()
+        assert self.model_path.exists()
+        self.n_threads = n_threads
+        if use_int8:
+            self.compute_type = "int8"
+        else:
+            self.compute_type = "default"
+        self._model = None
+
+    def init(self):
+        if self._model is None:
+            self.logger.debug("Initializing ...")
+            self._model = _WhisperModel(
+                self.model_path.__str__(),
+                device="cpu",
+                cpu_threads=self.n_threads,
+                compute_type=self.compute_type,
+                local_files_only=True)
+            self.logger.debug("Initialized.")
+
+    def transcribe(self, audio: bytes) -> str:
+        '''Transcibes audio and yields the segment strings'''
+        assert type(audio) == bytes
+        with _BytesIO(audio) as bio:
+            bio.seek(0)
+            self.init()
+            self.logger.debug("Transcribing audio ...")
+            segments, _ = self._model.transcribe(
+                bio,
+                beam_size=5, # beam size -> performance/quality
+                vad_filter=True) # remove silence
+            for s in segments:
+                self.logger.debug(f"... segment #{s.id}")
+                if s.id == 1:
+                    text = s.text.lstrip(" ")
+                else:
+                    text = s.text
+                yield text