b5f82fb48c
Some checks are pending
CI / ruff (push) Waiting to run
CI / mypy (push) Waiting to run
CI / pytest (push) Waiting to run
CI / package (push) Waiting to run
CI / Install smoke (${{ matrix.label }}) (linux, ubuntu-latest) (push) Blocked by required conditions
CI / Install smoke (${{ matrix.label }}) (macos-arm64, macos-14) (push) Blocked by required conditions
73 lines
1.9 KiB
Python
73 lines
1.9 KiB
Python
import os
|
|
|
|
os.environ['KERAS_BACKEND'] = 'torch'
|
|
|
|
import logging
|
|
|
|
import moonshine
|
|
import torch
|
|
from rich.console import Console
|
|
|
|
from speech_to_speech.baseHandler import BaseHandler
|
|
from speech_to_speech.pipeline.messages import VADAudio
|
|
|
|
logger = logging.getLogger(__name__)
|
|
console = Console()
|
|
|
|
|
|
class MoonshineSTTHandler(BaseHandler[VADAudio]):
|
|
"""
|
|
Handles the Speech To Text generation using a Moonshine model.
|
|
"""
|
|
|
|
def setup(
|
|
self,
|
|
model_name="moonshine/base",
|
|
torch_dtype="float16",
|
|
gen_kwargs={},
|
|
):
|
|
self.torch_dtype = getattr(torch, torch_dtype)
|
|
self.gen_kwargs = gen_kwargs
|
|
|
|
self.tokenizer = moonshine.load_tokenizer()
|
|
self.model = moonshine.load_model(model_name)
|
|
|
|
self.warmup()
|
|
|
|
def warmup(self):
|
|
logger.info(f"Warming up {self.__class__.__name__}")
|
|
|
|
n_steps = 2
|
|
dummy_input = torch.randn(
|
|
(1, 16000),
|
|
dtype=self.torch_dtype,
|
|
)
|
|
|
|
if torch.cuda.is_available():
|
|
start_event = torch.cuda.Event(enable_timing=True)
|
|
end_event = torch.cuda.Event(enable_timing=True)
|
|
torch.cuda.synchronize()
|
|
start_event.record()
|
|
|
|
for _ in range(n_steps):
|
|
_ = self.model.generate(dummy_input)
|
|
|
|
if torch.cuda.is_available():
|
|
end_event.record()
|
|
torch.cuda.synchronize()
|
|
|
|
logger.info(
|
|
f"{self.__class__.__name__}: warmed up! time: {start_event.elapsed_time(end_event) * 1e-3:.3f} s"
|
|
)
|
|
|
|
def process(self, vad_audio: VADAudio):
|
|
logger.debug("infering moonshine...")
|
|
|
|
pred_ids = self.model.generate(vad_audio.audio[None, :])
|
|
pred_text = self.tokenizer.decode_batch(pred_ids)[0]
|
|
|
|
logger.debug("finished whisper inference")
|
|
console.print(f"[yellow]USER: {pred_text}")
|
|
|
|
yield (pred_text, "en")
|