mirror of
https://github.com/HumanAIGC-Engineering/gradio-webrtc.git
synced 2026-02-04 17:39:23 +08:00
stt models (#147)
This commit is contained in:
@@ -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)
|
||||
|
||||
@@ -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"]
|
||||
|
||||
@@ -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
|
||||
]
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user