220 lines
8.4 KiB
Python
220 lines
8.4 KiB
Python
import sys
|
|
import numpy as np
|
|
import logging
|
|
from typing import List, Tuple, Optional
|
|
import logging
|
|
from whisperlivekit.timed_objects import ASRToken, Transcript
|
|
from whisperlivekit.simul_whisper.license_simulstreaming import SIMULSTREAMING_LICENSE
|
|
logger = logging.getLogger(__name__)
|
|
|
|
try:
|
|
import torch
|
|
from whisperlivekit.simul_whisper.config import AlignAttConfig
|
|
from whisperlivekit.simul_whisper.simul_whisper import PaddedAlignAttWhisper
|
|
from whisperlivekit.simul_whisper.whisper import tokenizer
|
|
except ImportError as e:
|
|
raise ImportError(
|
|
"""SimulStreaming dependencies are not available.
|
|
Please install WhisperLiveKit using pip install "whisperlivekit[simulstreaming]".""")
|
|
|
|
class SimulStreamingOnlineProcessor:
|
|
SAMPLING_RATE = 16000
|
|
|
|
def __init__(
|
|
self,
|
|
asr,
|
|
logfile=sys.stderr,
|
|
):
|
|
self.asr = asr
|
|
self.logfile = logfile
|
|
self.is_last = False
|
|
self.beg = 0.0
|
|
self.end = 0.0
|
|
self.cumulative_audio_duration = 0.0
|
|
|
|
self.committed: List[ASRToken] = []
|
|
self.last_result_tokens: List[ASRToken] = []
|
|
self.buffer_content = ""
|
|
|
|
def insert_audio_chunk(self, audio: np.ndarray, audio_stream_end_time: Optional[float] = None):
|
|
"""Append an audio chunk to be processed by SimulStreaming."""
|
|
|
|
# Convert numpy array to torch tensor
|
|
audio_tensor = torch.from_numpy(audio).float()
|
|
|
|
# Update timing
|
|
chunk_duration = len(audio) / self.SAMPLING_RATE
|
|
self.cumulative_audio_duration += chunk_duration
|
|
|
|
if audio_stream_end_time is not None:
|
|
self.end = audio_stream_end_time
|
|
else:
|
|
self.end = self.cumulative_audio_duration
|
|
self.asr.model.insert_audio(audio_tensor)
|
|
|
|
def get_buffer(self):
|
|
"""
|
|
Get the unvalidated buffer content.
|
|
"""
|
|
buffer_end = self.end if hasattr(self, 'end') else None
|
|
return Transcript(
|
|
start=None,
|
|
end=buffer_end,
|
|
text=self.buffer_content,
|
|
probability=None
|
|
)
|
|
|
|
def timestamped_text(self, tokens, generation):
|
|
# From the simulstreaming repo. self.model to self.asr.model
|
|
pr = generation["progress"]
|
|
if "result" not in generation:
|
|
split_words, split_tokens = self.asr.model.tokenizer.split_to_word_tokens(tokens)
|
|
else:
|
|
split_words, split_tokens = generation["result"]["split_words"], generation["result"]["split_tokens"]
|
|
|
|
frames = [p["most_attended_frames"][0] for p in pr]
|
|
tokens = tokens.copy()
|
|
ret = []
|
|
for sw,st in zip(split_words,split_tokens):
|
|
b = None
|
|
for stt in st:
|
|
t,f = tokens.pop(0), frames.pop(0)
|
|
if t != stt:
|
|
raise ValueError(f"Token mismatch: {t} != {stt} at frame {f}.")
|
|
if b is None:
|
|
b = f
|
|
e = f
|
|
out = (b*0.02, e*0.02, sw)
|
|
ret.append(out)
|
|
logger.debug(f"TS-WORD:\t{' '.join(map(str, out))}")
|
|
return ret
|
|
|
|
def process_iter(self) -> Tuple[List[ASRToken], float]:
|
|
"""
|
|
Process accumulated audio chunks using SimulStreaming.
|
|
|
|
Returns a tuple: (list of committed ASRToken objects, float representing the audio processed up to time).
|
|
"""
|
|
try:
|
|
tokens, generation_progress = self.asr.model.infer(is_last=self.is_last)
|
|
ts_words = self.timestamped_text(tokens, generation_progress)
|
|
|
|
new_tokens = []
|
|
for ts_word in ts_words:
|
|
|
|
start, end, word = ts_word
|
|
token = ASRToken(
|
|
start=start,
|
|
end=end,
|
|
text=word,
|
|
probability=0.95 # fake prob. Maybe we can extract it from the model?
|
|
)
|
|
new_tokens.append(token)
|
|
self.committed.extend(new_tokens)
|
|
|
|
return new_tokens, self.end
|
|
|
|
|
|
except Exception as e:
|
|
logger.exception(f"SimulStreaming processing error: {e}")
|
|
return [], self.end
|
|
|
|
class SimulStreamingASR():
|
|
"""SimulStreaming backend with AlignAtt policy."""
|
|
sep = ""
|
|
|
|
def __init__(self, lan, modelsize=None, cache_dir=None, model_dir=None, logfile=sys.stderr, **kwargs):
|
|
logger.warning(SIMULSTREAMING_LICENSE)
|
|
self.logfile = logfile
|
|
self.transcribe_kargs = {}
|
|
self.original_language = None if lan == "auto" else lan
|
|
|
|
self.model_path = kwargs.get('model_path', './large-v3.pt')
|
|
self.frame_threshold = kwargs.get('frame_threshold', 25)
|
|
self.audio_max_len = kwargs.get('audio_max_len', 30.0)
|
|
self.audio_min_len = kwargs.get('audio_min_len', 0.0)
|
|
self.segment_length = kwargs.get('segment_length', 0.5)
|
|
self.beams = kwargs.get('beams', 1)
|
|
self.decoder_type = kwargs.get('decoder_type', 'greedy' if self.beams == 1 else 'beam')
|
|
self.task = kwargs.get('task', 'transcribe')
|
|
self.cif_ckpt_path = kwargs.get('cif_ckpt_path', None)
|
|
self.never_fire = kwargs.get('never_fire', False)
|
|
self.init_prompt = kwargs.get('init_prompt', None)
|
|
self.static_init_prompt = kwargs.get('static_init_prompt', None)
|
|
self.max_context_tokens = kwargs.get('max_context_tokens', None)
|
|
|
|
if model_dir is not None:
|
|
self.model_path = model_dir
|
|
elif modelsize is not None:
|
|
model_mapping = {
|
|
'tiny': './tiny.pt',
|
|
'base': './base.pt',
|
|
'small': './small.pt',
|
|
'medium': './medium.pt',
|
|
'medium.en': './medium.en.pt',
|
|
'large-v1': './large-v1.pt',
|
|
'base.en': './base.en.pt',
|
|
'small.en': './small.en.pt',
|
|
'tiny.en': './tiny.en.pt',
|
|
'large-v2': './large-v2.pt',
|
|
'large-v3': './large-v3.pt',
|
|
'large': './large-v3.pt'
|
|
}
|
|
self.model_path = model_mapping.get(modelsize, f'./{modelsize}.pt')
|
|
|
|
self.model = self.load_model(modelsize, cache_dir, model_dir)
|
|
|
|
# Set up tokenizer for translation if needed
|
|
if self.task == "translate":
|
|
self.set_translate_task()
|
|
|
|
def load_model(self, modelsize, cache_dir, model_dir):
|
|
try:
|
|
cfg = AlignAttConfig(
|
|
model_path=self.model_path,
|
|
segment_length=self.segment_length,
|
|
frame_threshold=self.frame_threshold,
|
|
language=self.original_language,
|
|
audio_max_len=self.audio_max_len,
|
|
audio_min_len=self.audio_min_len,
|
|
cif_ckpt_path=self.cif_ckpt_path,
|
|
decoder_type="beam",
|
|
beam_size=self.beams,
|
|
task=self.task,
|
|
never_fire=self.never_fire,
|
|
init_prompt=self.init_prompt,
|
|
max_context_tokens=self.max_context_tokens,
|
|
static_init_prompt=self.static_init_prompt,
|
|
)
|
|
model = PaddedAlignAttWhisper(cfg)
|
|
return model
|
|
|
|
except Exception as e:
|
|
logger.error(f"Failed to load SimulStreaming model: {e}")
|
|
raise
|
|
|
|
def set_translate_task(self):
|
|
"""Set up translation task."""
|
|
try:
|
|
self.model.tokenizer = tokenizer.get_tokenizer(
|
|
multilingual=True,
|
|
language=self.model.cfg.language,
|
|
num_languages=self.model.model.num_languages,
|
|
task="translate"
|
|
)
|
|
logger.info("SimulStreaming configured for translation task")
|
|
except Exception as e:
|
|
logger.error(f"Failed to configure SimulStreaming for translation: {e}")
|
|
raise
|
|
|
|
def warmup(self, audio, init_prompt=""):
|
|
"""Warmup the SimulStreaming model."""
|
|
try:
|
|
if isinstance(audio, np.ndarray):
|
|
audio = torch.from_numpy(audio).float()
|
|
self.model.insert_audio(audio)
|
|
self.model.infer(True)
|
|
self.model.refresh_segment(complete=True)
|
|
logger.info("SimulStreaming model warmed up successfully")
|
|
except Exception as e:
|
|
logger.exception(f"SimulStreaming warmup failed: {e}")
|