57 lines
2 KiB
Python
57 lines
2 KiB
Python
|
# 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
|