Pine
7f25c8f6bb
- main.py：uv run main.py 一键启动——读取 backend/.env、TORCH_HOME/NLTK_DATA 指向项目内缓存、certifi SSL + NO_PROXY 网络直连、默认开启 s2s，启动横幅 - vendor/：补丁后 s2s 云化栈移入项目（backend/vendor/s2s-cloud），pyproject sources 改相对路径 - .gitignore：忽略 .torch-cache / nltk_data 缓存 - AGENTS.md：进程管理铁律（AI 不启动/重启服务，由用户操作）
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")
|