stt models (#147)

This commit is contained in:
Freddy Boulton
2025-03-07 17:03:11 -05:00
committed by GitHub
parent cbbfa17679
commit 504eb452f0
6 changed files with 55 additions and 64 deletions

View File

@@ -13,7 +13,7 @@ from .reply_on_pause import (
ReplyFnGenerator,
ReplyOnPause,
)
from .speech_to_text import get_stt_model
from .speech_to_text import get_stt_model, stt_for_chunks
from .utils import audio_to_float32, create_message
logger = logging.getLogger(__name__)
@@ -105,10 +105,9 @@ class ReplyOnStopWords(ReplyOnPause):
dur_vad, chunks = self.model.vad(
(16000, state.post_stop_word_buffer),
self.model_options,
return_chunks=True,
)
text = self.stt_model.stt_for_chunks(
(16000, state.post_stop_word_buffer), chunks
text = stt_for_chunks(
self.stt_model, (16000, state.post_stop_word_buffer), chunks
)
logger.debug(f"STT: {text}")
state.stop_word_detected = self.stop_word_detected(text)

View File

@@ -1,3 +1,3 @@
from .stt_ import MoonshineSTT, get_stt_model
from .stt_ import MoonshineSTT, get_stt_model, stt_for_chunks
__all__ = ["get_stt_model", "MoonshineSTT", "get_stt_model"]
__all__ = ["get_stt_model", "MoonshineSTT", "get_stt_model", "stt_for_chunks"]

View File

@@ -15,12 +15,6 @@ curr_dir = Path(__file__).parent
class STTModel(Protocol):
def stt(self, audio: tuple[int, NDArray[np.int16 | np.float32]]) -> str: ...
def stt_for_chunks(
self,
audio: tuple[int, NDArray[np.int16 | np.float32]],
chunks: list[AudioChunk],
) -> str: ...
class MoonshineSTT(STTModel):
def __init__(
@@ -49,19 +43,6 @@ class MoonshineSTT(STTModel):
tokens = self.model.generate(audio_np)
return self.tokenizer.decode_batch(tokens)[0]
def stt_for_chunks(
self,
audio: tuple[int, NDArray[np.int16 | np.float32]],
chunks: list[AudioChunk],
) -> str:
sr, audio_np = audio
return " ".join(
[
self.stt((sr, audio_np[chunk["start"] : chunk["end"]]))
for chunk in chunks
]
)
@lru_cache
def get_stt_model(
@@ -79,3 +60,17 @@ def get_stt_model(
m.stt((16000, audio))
print(click.style("INFO", fg="green") + ":\t STT model warmed up.")
return m
def stt_for_chunks(
stt_model: STTModel,
audio: tuple[int, NDArray[np.int16 | np.float32]],
chunks: list[AudioChunk],
) -> str:
sr, audio_np = audio
return " ".join(
[
stt_model.stt((sr, audio_np[chunk["start"] : chunk["end"]]))
for chunk in chunks
]
)