mirror of
https://github.com/QuentinFuxa/WhisperLiveKit.git
synced 2026-03-07 22:33:36 +00:00
445 lines
18 KiB
Python
445 lines
18 KiB
Python
import sys
|
|
import numpy as np
|
|
import logging
|
|
from typing import List, Tuple, Optional
|
|
from src.whisper_streaming.timed_objects import ASRToken, Sentence, Transcript
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class HypothesisBuffer:
|
|
"""
|
|
Buffer to store and process ASR hypothesis tokens.
|
|
|
|
It holds:
|
|
- committed_in_buffer: tokens that have been confirmed (committed)
|
|
- buffer: the last hypothesis that is not yet committed
|
|
- new: new tokens coming from the recognizer
|
|
"""
|
|
def __init__(self, logfile=sys.stderr):
|
|
self.committed_in_buffer: List[ASRToken] = []
|
|
self.buffer: List[ASRToken] = []
|
|
self.new: List[ASRToken] = []
|
|
self.last_committed_time = 0.0
|
|
self.last_committed_word: Optional[str] = None
|
|
self.logfile = logfile
|
|
|
|
def insert(self, new_tokens: List[ASRToken], offset: float):
|
|
"""
|
|
Insert new tokens (after applying a time offset) and compare them with the
|
|
already committed tokens. Only tokens that extend the committed hypothesis
|
|
are added.
|
|
"""
|
|
# Apply the offset to each token.
|
|
new_tokens = [token.with_offset(offset) for token in new_tokens]
|
|
# Only keep tokens that are roughly “new”
|
|
self.new = [token for token in new_tokens if token.start > self.last_committed_time - 0.1]
|
|
|
|
if self.new:
|
|
first_token = self.new[0]
|
|
if abs(first_token.start - self.last_committed_time) < 1:
|
|
if self.committed_in_buffer:
|
|
committed_len = len(self.committed_in_buffer)
|
|
new_len = len(self.new)
|
|
# Try to match 1 to 5 consecutive tokens
|
|
max_ngram = min(min(committed_len, new_len), 5)
|
|
for i in range(1, max_ngram + 1):
|
|
committed_ngram = " ".join(token.text for token in self.committed_in_buffer[-i:])
|
|
new_ngram = " ".join(token.text for token in self.new[:i])
|
|
if committed_ngram == new_ngram:
|
|
removed = []
|
|
for _ in range(i):
|
|
removed_token = self.new.pop(0)
|
|
removed.append(repr(removed_token))
|
|
logger.debug(f"Removing last {i} words: {' '.join(removed)}")
|
|
break
|
|
|
|
def flush(self) -> List[ASRToken]:
|
|
"""
|
|
Returns the committed chunk, defined as the longest common prefix
|
|
between the previous hypothesis and the new tokens.
|
|
"""
|
|
committed: List[ASRToken] = []
|
|
while self.new:
|
|
current_new = self.new[0]
|
|
if not self.buffer:
|
|
break
|
|
if current_new.text == self.buffer[0].text:
|
|
committed.append(current_new)
|
|
self.last_committed_word = current_new.text
|
|
self.last_committed_time = current_new.end
|
|
self.buffer.pop(0)
|
|
self.new.pop(0)
|
|
else:
|
|
break
|
|
self.buffer = self.new
|
|
self.new = []
|
|
self.committed_in_buffer.extend(committed)
|
|
return committed
|
|
|
|
def pop_committed(self, time: float):
|
|
"""
|
|
Remove tokens (from the beginning) that have ended before `time`.
|
|
"""
|
|
while self.committed_in_buffer and self.committed_in_buffer[0].end <= time:
|
|
self.committed_in_buffer.pop(0)
|
|
|
|
|
|
|
|
class OnlineASRProcessor:
|
|
"""
|
|
Processes incoming audio in a streaming fashion, calling the ASR system
|
|
periodically, and uses a hypothesis buffer to commit and trim recognized text.
|
|
|
|
The processor supports two types of buffer trimming:
|
|
- "sentence": trims at sentence boundaries (using a sentence tokenizer)
|
|
- "segment": trims at fixed segment durations.
|
|
"""
|
|
SAMPLING_RATE = 16000
|
|
|
|
def __init__(
|
|
self,
|
|
asr,
|
|
tokenize_method: Optional[callable] = None,
|
|
buffer_trimming: Tuple[str, float] = ("segment", 15),
|
|
logfile=sys.stderr,
|
|
):
|
|
"""
|
|
asr: An ASR system object (for example, a WhisperASR instance) that
|
|
provides a `transcribe` method, a `ts_words` method (to extract tokens),
|
|
a `segments_end_ts` method, and a separator attribute `sep`.
|
|
tokenize_method: A function that receives text and returns a list of sentence strings.
|
|
buffer_trimming: A tuple (option, seconds), where option is either "sentence" or "segment".
|
|
"""
|
|
self.asr = asr
|
|
self.tokenize = tokenize_method
|
|
self.logfile = logfile
|
|
|
|
self.init()
|
|
|
|
self.buffer_trimming_way, self.buffer_trimming_sec = buffer_trimming
|
|
|
|
if self.buffer_trimming_way not in ["sentence", "segment"]:
|
|
raise ValueError("buffer_trimming must be either 'sentence' or 'segment'")
|
|
if self.buffer_trimming_sec <= 0:
|
|
raise ValueError("buffer_trimming_sec must be positive")
|
|
elif self.buffer_trimming_sec > 30:
|
|
logger.warning(
|
|
f"buffer_trimming_sec is set to {self.buffer_trimming_sec}, which is very long. It may cause OOM."
|
|
)
|
|
|
|
def init(self, offset: Optional[float] = None):
|
|
"""Initialize or reset the processing buffers."""
|
|
self.audio_buffer = np.array([], dtype=np.float32)
|
|
self.transcript_buffer = HypothesisBuffer(logfile=self.logfile)
|
|
self.buffer_time_offset = offset if offset is not None else 0.0
|
|
self.transcript_buffer.last_committed_time = self.buffer_time_offset
|
|
self.committed: List[ASRToken] = []
|
|
|
|
def insert_audio_chunk(self, audio: np.ndarray):
|
|
"""Append an audio chunk (a numpy array) to the current audio buffer."""
|
|
self.audio_buffer = np.append(self.audio_buffer, audio)
|
|
|
|
def prompt(self) -> Tuple[str, str]:
|
|
"""
|
|
Returns a tuple: (prompt, context), where:
|
|
- prompt is a 200-character suffix of committed text that falls
|
|
outside the current audio buffer.
|
|
- context is the committed text within the current audio buffer.
|
|
"""
|
|
k = len(self.committed)
|
|
while k > 0 and self.committed[k - 1].end > self.buffer_time_offset:
|
|
k -= 1
|
|
|
|
prompt_tokens = self.committed[:k]
|
|
prompt_words = [token.text for token in prompt_tokens]
|
|
prompt_list = []
|
|
length_count = 0
|
|
# Use the last words until reaching 200 characters.
|
|
while prompt_words and length_count < 200:
|
|
word = prompt_words.pop(-1)
|
|
length_count += len(word) + 1
|
|
prompt_list.append(word)
|
|
non_prompt_tokens = self.committed[k:]
|
|
context_text = self.asr.sep.join(token.text for token in non_prompt_tokens)
|
|
return self.asr.sep.join(prompt_list[::-1]), context_text
|
|
|
|
def get_buffer(self):
|
|
"""
|
|
Get the unvalidated buffer in string format.
|
|
"""
|
|
return self.concatenate_tokens(self.transcript_buffer.buffer).text
|
|
|
|
|
|
def process_iter(self) -> Transcript:
|
|
"""
|
|
Processes the current audio buffer.
|
|
|
|
Returns a Transcript object representing the committed transcript.
|
|
"""
|
|
prompt_text, _ = self.prompt()
|
|
logger.debug(
|
|
f"Transcribing {len(self.audio_buffer)/self.SAMPLING_RATE:.2f} seconds from {self.buffer_time_offset:.2f}"
|
|
)
|
|
res = self.asr.transcribe(self.audio_buffer, init_prompt=prompt_text)
|
|
tokens = self.asr.ts_words(res) # Expecting List[ASRToken]
|
|
self.transcript_buffer.insert(tokens, self.buffer_time_offset)
|
|
committed_tokens = self.transcript_buffer.flush()
|
|
self.committed.extend(committed_tokens)
|
|
completed = self.concatenate_tokens(committed_tokens)
|
|
logger.debug(f">>>> COMPLETE NOW: {completed.text}")
|
|
incomp = self.concatenate_tokens(self.transcript_buffer.buffer)
|
|
logger.debug(f"INCOMPLETE: {incomp.text}")
|
|
|
|
if committed_tokens and self.buffer_trimming_way == "sentence":
|
|
if len(self.audio_buffer) / self.SAMPLING_RATE > self.buffer_trimming_sec:
|
|
self.chunk_completed_sentence()
|
|
|
|
s = self.buffer_trimming_sec if self.buffer_trimming_way == "segment" else 30
|
|
if len(self.audio_buffer) / self.SAMPLING_RATE > s:
|
|
self.chunk_completed_segment(res)
|
|
logger.debug("Chunking segment")
|
|
logger.debug(
|
|
f"Length of audio buffer now: {len(self.audio_buffer)/self.SAMPLING_RATE:.2f} seconds"
|
|
)
|
|
return self.concatenate_tokens(committed_tokens)
|
|
|
|
def chunk_completed_sentence(self):
|
|
"""
|
|
If the committed tokens form at least two sentences, chunk the audio
|
|
buffer at the end time of the penultimate sentence.
|
|
"""
|
|
if not self.committed:
|
|
return
|
|
logger.debug("COMPLETED SENTENCE: " + " ".join(token.text for token in self.committed))
|
|
sentences = self.words_to_sentences(self.committed)
|
|
for sentence in sentences:
|
|
logger.debug(f"\tSentence: {sentence.text}")
|
|
if len(sentences) < 2:
|
|
return
|
|
# Keep the last two sentences.
|
|
while len(sentences) > 2:
|
|
sentences.pop(0)
|
|
chunk_time = sentences[-2].end
|
|
logger.debug(f"--- Sentence chunked at {chunk_time:.2f}")
|
|
self.chunk_at(chunk_time)
|
|
|
|
def chunk_completed_segment(self, res):
|
|
"""
|
|
Chunk the audio buffer based on segment-end timestamps reported by the ASR.
|
|
"""
|
|
if not self.committed:
|
|
return
|
|
ends = self.asr.segments_end_ts(res)
|
|
last_committed_time = self.committed[-1].end
|
|
if len(ends) > 1:
|
|
e = ends[-2] + self.buffer_time_offset
|
|
while len(ends) > 2 and e > last_committed_time:
|
|
ends.pop(-1)
|
|
e = ends[-2] + self.buffer_time_offset
|
|
if e <= last_committed_time:
|
|
logger.debug(f"--- Segment chunked at {e:.2f}")
|
|
self.chunk_at(e)
|
|
else:
|
|
logger.debug("--- Last segment not within committed area")
|
|
else:
|
|
logger.debug("--- Not enough segments to chunk")
|
|
|
|
def chunk_at(self, time: float):
|
|
"""
|
|
Trim both the hypothesis and audio buffer at the given time.
|
|
"""
|
|
logger.debug(f"Chunking at {time:.2f}s")
|
|
logger.debug(
|
|
f"Audio buffer length before chunking: {len(self.audio_buffer)/self.SAMPLING_RATE:.2f}s"
|
|
)
|
|
self.transcript_buffer.pop_committed(time)
|
|
cut_seconds = time - self.buffer_time_offset
|
|
self.audio_buffer = self.audio_buffer[int(cut_seconds * self.SAMPLING_RATE):]
|
|
self.buffer_time_offset = time
|
|
logger.debug(
|
|
f"Audio buffer length after chunking: {len(self.audio_buffer)/self.SAMPLING_RATE:.2f}s"
|
|
)
|
|
|
|
def words_to_sentences(self, tokens: List[ASRToken]) -> List[Sentence]:
|
|
"""
|
|
Converts a list of tokens to a list of Sentence objects using the provided
|
|
sentence tokenizer.
|
|
"""
|
|
if not tokens:
|
|
return []
|
|
|
|
full_text = " ".join(token.text for token in tokens)
|
|
|
|
if self.tokenize:
|
|
try:
|
|
sentence_texts = self.tokenize(full_text)
|
|
except Exception as e:
|
|
# Some tokenizers (e.g., MosesSentenceSplitter) expect a list input.
|
|
try:
|
|
sentence_texts = self.tokenize([full_text])
|
|
except Exception as e2:
|
|
raise ValueError("Tokenization failed") from e2
|
|
else:
|
|
sentence_texts = [full_text]
|
|
|
|
sentences: List[Sentence] = []
|
|
token_index = 0
|
|
for sent_text in sentence_texts:
|
|
sent_text = sent_text.strip()
|
|
if not sent_text:
|
|
continue
|
|
sent_tokens = []
|
|
accumulated = ""
|
|
# Accumulate tokens until roughly matching the length of the sentence text.
|
|
while token_index < len(tokens) and len(accumulated) < len(sent_text):
|
|
token = tokens[token_index]
|
|
accumulated = (accumulated + " " + token.text).strip() if accumulated else token.text
|
|
sent_tokens.append(token)
|
|
token_index += 1
|
|
if sent_tokens:
|
|
sentence = Sentence(
|
|
start=sent_tokens[0].start,
|
|
end=sent_tokens[-1].end,
|
|
text=" ".join(t.text for t in sent_tokens),
|
|
)
|
|
sentences.append(sentence)
|
|
return sentences
|
|
def finish(self) -> Transcript:
|
|
"""
|
|
Flush the remaining transcript when processing ends.
|
|
"""
|
|
remaining_tokens = self.transcript_buffer.buffer
|
|
final_transcript = self.concatenate_tokens(remaining_tokens)
|
|
logger.debug(f"Final non-committed transcript: {final_transcript}")
|
|
self.buffer_time_offset += len(self.audio_buffer) / self.SAMPLING_RATE
|
|
return final_transcript
|
|
|
|
def concatenate_tokens(
|
|
self,
|
|
tokens: List[ASRToken],
|
|
sep: Optional[str] = None,
|
|
offset: float = 0
|
|
) -> Transcript:
|
|
sep = sep if sep is not None else self.asr.sep
|
|
text = sep.join(token.text for token in tokens)
|
|
if tokens:
|
|
start = offset + tokens[0].start
|
|
end = offset + tokens[-1].end
|
|
else:
|
|
start = None
|
|
end = None
|
|
return Transcript(start, end, text)
|
|
|
|
|
|
class VACOnlineASRProcessor:
|
|
"""
|
|
Wraps an OnlineASRProcessor with a Voice Activity Controller (VAC).
|
|
|
|
It receives small chunks of audio, applies VAD (e.g. with Silero),
|
|
and when the system detects a pause in speech (or end of an utterance)
|
|
it finalizes the utterance immediately.
|
|
"""
|
|
SAMPLING_RATE = 16000
|
|
|
|
def __init__(self, online_chunk_size: float, *args, **kwargs):
|
|
self.online_chunk_size = online_chunk_size
|
|
self.online = OnlineASRProcessor(*args, **kwargs)
|
|
|
|
# Load a VAD model (e.g. Silero VAD)
|
|
import torch
|
|
model, _ = torch.hub.load(repo_or_dir="snakers4/silero-vad", model="silero_vad")
|
|
from src.whisper_streaming.silero_vad_iterator import FixedVADIterator
|
|
|
|
self.vac = FixedVADIterator(model)
|
|
self.logfile = self.online.logfile
|
|
self.init()
|
|
|
|
def init(self):
|
|
self.online.init()
|
|
self.vac.reset_states()
|
|
self.current_online_chunk_buffer_size = 0
|
|
self.is_currently_final = False
|
|
self.status: Optional[str] = None # "voice" or "nonvoice"
|
|
self.audio_buffer = np.array([], dtype=np.float32)
|
|
self.buffer_offset = 0 # in frames
|
|
|
|
def clear_buffer(self):
|
|
self.buffer_offset += len(self.audio_buffer)
|
|
self.audio_buffer = np.array([], dtype=np.float32)
|
|
|
|
def insert_audio_chunk(self, audio: np.ndarray):
|
|
"""
|
|
Process an incoming small audio chunk:
|
|
- run VAD on the chunk,
|
|
- decide whether to send the audio to the online ASR processor immediately,
|
|
- and/or to mark the current utterance as finished.
|
|
"""
|
|
res = self.vac(audio)
|
|
self.audio_buffer = np.append(self.audio_buffer, audio)
|
|
|
|
if res is not None:
|
|
# VAD returned a result; adjust the frame number
|
|
frame = list(res.values())[0] - self.buffer_offset
|
|
if "start" in res and "end" not in res:
|
|
self.status = "voice"
|
|
send_audio = self.audio_buffer[frame:]
|
|
self.online.init(offset=(frame + self.buffer_offset) / self.SAMPLING_RATE)
|
|
self.online.insert_audio_chunk(send_audio)
|
|
self.current_online_chunk_buffer_size += len(send_audio)
|
|
self.clear_buffer()
|
|
elif "end" in res and "start" not in res:
|
|
self.status = "nonvoice"
|
|
send_audio = self.audio_buffer[:frame]
|
|
self.online.insert_audio_chunk(send_audio)
|
|
self.current_online_chunk_buffer_size += len(send_audio)
|
|
self.is_currently_final = True
|
|
self.clear_buffer()
|
|
else:
|
|
beg = res["start"] - self.buffer_offset
|
|
end = res["end"] - self.buffer_offset
|
|
self.status = "nonvoice"
|
|
send_audio = self.audio_buffer[beg:end]
|
|
self.online.init(offset=(beg + self.buffer_offset) / self.SAMPLING_RATE)
|
|
self.online.insert_audio_chunk(send_audio)
|
|
self.current_online_chunk_buffer_size += len(send_audio)
|
|
self.is_currently_final = True
|
|
self.clear_buffer()
|
|
else:
|
|
if self.status == "voice":
|
|
self.online.insert_audio_chunk(self.audio_buffer)
|
|
self.current_online_chunk_buffer_size += len(self.audio_buffer)
|
|
self.clear_buffer()
|
|
else:
|
|
# Keep 1 second worth of audio in case VAD later detects voice,
|
|
# but trim to avoid unbounded memory usage.
|
|
self.buffer_offset += max(0, len(self.audio_buffer) - self.SAMPLING_RATE)
|
|
self.audio_buffer = self.audio_buffer[-self.SAMPLING_RATE:]
|
|
|
|
def process_iter(self) -> Transcript:
|
|
"""
|
|
Depending on the VAD status and the amount of accumulated audio,
|
|
process the current audio chunk.
|
|
"""
|
|
if self.is_currently_final:
|
|
return self.finish()
|
|
elif self.current_online_chunk_buffer_size > self.SAMPLING_RATE * self.online_chunk_size:
|
|
self.current_online_chunk_buffer_size = 0
|
|
return self.online.process_iter()
|
|
else:
|
|
logger.debug("No online update, only VAD")
|
|
return Transcript(None, None, "")
|
|
|
|
def finish(self) -> Transcript:
|
|
"""Finish processing by flushing any remaining text."""
|
|
result = self.online.finish()
|
|
self.current_online_chunk_buffer_size = 0
|
|
self.is_currently_final = False
|
|
return result
|
|
|
|
def get_buffer(self):
|
|
"""
|
|
Get the unvalidated buffer in string format.
|
|
"""
|
|
return self.online.concatenate_tokens(self.online.transcript_buffer.buffer).text
|