69 lines
2.1 KiB
Python
69 lines
2.1 KiB
Python
from __future__ import annotations
|
|
|
|
import logging
|
|
import os
|
|
from typing import Any, Iterator
|
|
|
|
from faster_whisper import WhisperModel
|
|
from rich.console import Console
|
|
|
|
from speech_to_speech.pipeline.handler_types import STTIn, STTOut
|
|
from speech_to_speech.pipeline.messages import Transcription
|
|
from speech_to_speech.STT.base_stt_handler import BaseSTTHandler
|
|
|
|
console = Console()
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class FasterWhisperSTTHandler(BaseSTTHandler):
|
|
"""
|
|
Handles the Speech To Text generation using a Whisper model.
|
|
"""
|
|
|
|
def setup(
|
|
self,
|
|
model_name: str = "tiny.en",
|
|
device: str = "auto",
|
|
compute_type: str = "auto",
|
|
gen_kwargs: dict[str, Any] = {},
|
|
) -> None:
|
|
self.gen_kwargs = self.adapt_gen_kwargs(gen_kwargs)
|
|
|
|
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
|
|
self.model = WhisperModel(model_name, device=device, compute_type=compute_type)
|
|
|
|
def process(self, vad_audio: STTIn) -> Iterator[STTOut]:
|
|
logger.debug("infering faster whisper...")
|
|
|
|
segments, info = self.model.transcribe(vad_audio.audio, **self.gen_kwargs)
|
|
output_text = []
|
|
|
|
for segment in segments:
|
|
logger.debug("[%.2fs -> %.2fs] %s" % (segment.start, segment.end, segment.text))
|
|
output_text.append(segment.text)
|
|
|
|
pred_text = " ".join(output_text).strip()
|
|
|
|
logger.debug("finished whisper inference")
|
|
if pred_text:
|
|
console.print(f"[yellow]USER: {pred_text}")
|
|
|
|
yield Transcription(
|
|
text=pred_text,
|
|
turn_id=vad_audio.turn_id,
|
|
turn_revision=vad_audio.turn_revision,
|
|
speech_stopped_at_s=vad_audio.created_at_s,
|
|
)
|
|
else:
|
|
logger.debug("no text detected. skipping...")
|
|
|
|
def cleanup(self) -> None:
|
|
print("Stopping FasterWhisperSTTHandler")
|
|
del self.model
|
|
|
|
def adapt_gen_kwargs(self, gen_kwargs: dict[str, Any]) -> dict[str, Any]:
|
|
gen_kwargs["without_timestamps"] = not gen_kwargs.pop("return_timestamps", True)
|
|
|
|
return gen_kwargs
|