import numpy as np

import threading

from loguru import logger

import torch

import whisperx

SAMPLING_RATE = 16000 # Hz

MIN_AUDIO_BUFFER_DURATION = 1 # seconds

MAX_AUDIO_BUFFER_DURATION = 15 # seconds

FAST_FORWARD_TIME_MARGIN = 0.1 # seconds

ASR_CONTEXT_LENGTH = 200 # words

VAD_SILENCE_STT_THRESHOLD = 0.2

CONDITIONING_TEXT = {

"en": "He thoughtfuly said:",

"ru": "Он задумчиво сказал:"

}

cuda_lock = threading.Lock()

class AudioBuffer:

def __init__(

self,

sampling_rate=SAMPLING_RATE,

min_duration=MIN_AUDIO_BUFFER_DURATION,

max_duration=MAX_AUDIO_BUFFER_DURATION

):

self.sampling_rate = sampling_rate

self.min_duration = min_duration

self.max_duration = max_duration

self.min_size = int(min_duration * sampling_rate)

self.max_size = int(max_duration * sampling_rate)

self.buffer = None

self.offset = None # global time of the buffer's first element

self.clear()

def push(self, data):

if data is not None:

assert isinstance(data, np.ndarray), f"Invalid data type: {type(data)}"

self.buffer = np.concatenate([self.buffer, data])

if len(self.buffer) < self.min_size:

logger.debug("Audio buffer ({:.2f}s) is shorter than {}s", len(self.buffer) / self.sampling_rate, self.min_duration)

return None, None

elif len(self.buffer) > self.max_size:

o_offset = self.offset

logger.debug("Audio buffer ({:.2f}s) is longer than {}s, trimming.", len(self.buffer) / self.sampling_rate, self.max_duration)

self.offset += (len(self.buffer) - self.max_size) / self.sampling_rate

self.buffer = self.buffer[-self.max_size:]

logger.debug("New buffer length: {:.2f}s. Offset: {:.2f}s -> {:.2f}s", len(self.buffer) / self.sampling_rate, o_offset, self.offset)

if data is None:

return self.clear()

return self.buffer, self.offset

def clear(self):

buffer, offset = self.buffer, self.offset

self.buffer = np.empty((0,), dtype=np.float32)

self.offset = 0

return buffer, offset

def fast_forward(self, time):

if time < self.offset:

logger.debug("Ignoring negative fast forward: {:.2f} -> {:.2f}s. Maybe the buffer was front-trimmed?", self.offset, time)

return

ff = min(time - self.offset, len(self.buffer) / self.sampling_rate)

ff_size = int(ff * self.sampling_rate)

self.buffer = self.buffer[ff_size:]

self.offset += ff

def empty(self):

return len(self.buffer) == 0

class Word:

def __init__(self, word, start, end, sep=""):

self.word = word

self.start = start

self.end = end

self.sep = sep

@staticmethod

def none():

return Word(word=None, start=None, end=None)

@staticmethod

def to_text(words):

words = words if isinstance(words, list) else [words]

if not words:

return ""

return words[0].sep.join([w.word for w in words])

@staticmethod

def apply_offset(words, offset):

words = words if isinstance(words, list) else [words]

for i in range(len(words)):

words[i].start += offset

words[i].end += offset

return words

def __repr__(self):

return f"{self.start}:{self.end} {self.word}"

class HypothesisBuffer:

def __init__(self):

self.clear()

def clear(self):

self.confirmed_words = []

self.unconfirmed_words = []

def update(self, current_words):

if self.confirmed_words:

current_words = self._fast_forward(current_words, self.confirmed_words[-1].end)

self.unconfirmed_words = self._fast_forward(self.unconfirmed_words, self.confirmed_words[-1].end)

if self.unconfirmed_words:

lcp = self._longest_common_prefix(self.unconfirmed_words, current_words)

self.confirmed_words.extend(lcp)

if lcp:

current_words = self._fast_forward(current_words, self.confirmed_words[-1].end)

self.unconfirmed_words = current_words

return self.confirmed_words, self.unconfirmed_words

@staticmethod

def _longest_common_prefix(s1, s2):

min_len = min(len(s1), len(s2))

logger.debug("comparing:")

logger.debug("unconf: {}", Word.to_text(s1))

logger.debug("currnt: {}", Word.to_text(s2))

def compare(w1, w2):

w1_lower = ''.join(c for c in w1.word.lower() if c.isalpha())

w2_lower = ''.join(c for c in w2.word.lower() if c.isalpha())

return w1_lower == w2_lower

index = 0

while index < min_len and compare(s1[index], s2[index]):

index += 1

return s1[:index]

@staticmethod

def _fast_forward(words, time):

# fast forward `words` till the first word which starts after `time`

while words and words[0].start < time:

words.pop(0)

return words

class VoiceActivityDetector:

def __init__(self, sampling_rate=16000, voice_threshold=0.5):

assert sampling_rate in [8000, 16000], f"Unsupported sampling rate {sampling_rate}"

self.model, _ = torch.hub.load(repo_or_dir='snakers4/silero-vad', model='silero_vad')

self.sampling_rate = sampling_rate

self.voice_threshold = voice_threshold

self.silence_threshold = max(self.voice_threshold - 0.15, 0.1)

self._sample_width = {8000: 256, 16000: 512}[self.sampling_rate]

self._trailing_silence = 0.0

self._trailing_voice = 0.0

self._has_voice = None

self._remainder = None

self.reset()

def reset(self):

trailing_silence, trailing_voice, has_voice = (

self._trailing_silence, self._trailing_voice, self._has_voice

)

self._trailing_silence = 0.0

self._trailing_voice = 0.0

self._has_voice = False

self._remainder = np.empty((0,), dtype=np.float32)

return {

"trailing_silence": trailing_silence,

"trailing_voice": trailing_voice,

"has_voice": has_voice

}

def clear_voice(self):

self._trailing_voice = 0.0

self._has_voice = False

def process_chunk(self, chunk):

if chunk is None:

return self.reset()

chunk = np.concatenate([self._remainder, chunk])

for i in range(len(chunk) // self._sample_width):

sample = chunk[i * self._sample_width:(i + 1) * self._sample_width]

self._process_sample(sample)

self._remainder = chunk[-(len(chunk) % self._sample_width):]

return {

"trailing_silence": self._trailing_silence,

"trailing_voice": self._trailing_voice,

"has_voice": self._has_voice

}

def _process_sample(self, chunk):

prob = self.model(torch.from_numpy(chunk), self.sampling_rate)

if prob > self.voice_threshold:

self._trailing_voice += self._get_duration(chunk)

self._trailing_silence = 0.0

self._has_voice = True

elif prob < self.silence_threshold:

self._trailing_voice = 0.0

self._trailing_silence += self._get_duration(chunk)

def _get_duration(self, chunk):

return len(chunk) / self.sampling_rate

class OfflineASR:

word_sep = " "

_cached_model_params = None

@staticmethod

def get_model(language="en", device="cuda", cached=True):

if (cached and OfflineASR._cached_model_params and

OfflineASR._cached_model_params["language"] == language):

return OfflineASR._cached_model_params

if OfflineASR._cached_model_params:

OfflineASR._cached_model_params["model"].model.model.unload_model(to_cpu=True)

OfflineASR._cached_model_params["model_a"].cpu()

del OfflineASR._cached_model_params["model"]

del OfflineASR._cached_model_params["model_a"]

torch.cuda.empty_cache()

whisper_model = whisperx.load_model(

"large-v2",

device=device,

compute_type="float16",

language=language,

asr_options={"suppress_numerals": True}

# vad_options={'vad_onset': 0.8, 'vad_offset': 0.8}

)

align_model, align_metadata = whisperx.load_align_model(

language_code=language,

device=device,

model_name="WAV2VEC2_ASR_LARGE_LV60K_960H" if language == "en" else None

)

OfflineASR._cached_model_params = {

"model": whisper_model,

"model_a": align_model,

"align_metadata": align_metadata,

"language": language

}

return OfflineASR._cached_model_params

def __init__(self, language="en", cached=False):

self.language = language

self.batch_size = 1 # reduce if low on GPU mem

self.device = "cuda"

model_params = OfflineASR.get_model(language=language, cached=cached)

self.model = model_params["model"]

self.model_a = model_params["model_a"]

self.align_metadata = model_params["align_metadata"]

def transcribe(self, audio, previous_text=None):

with cuda_lock:

transcript = self.model.transcribe(

audio,

batch_size=self.batch_size,

language=self.language,

previous_text=previous_text

)

with cuda_lock:

result = whisperx.align(

transcript=transcript["segments"],

model=self.model_a,

align_model_metadata=self.align_metadata,

audio=audio,

device=self.device,

return_char_alignments=False

)

return self._to_words(result)

def _to_words(self, result):

words = []

for w in result["word_segments"]:

if "start" in w and "end" in w: # ignore unaligned words

word = Word(

word=w["word"],

start=w["start"],

end=w["end"],

sep=self.word_sep

)

words.append(word)

else:

logger.warning("Skipping unaligned word: {}", w)

return words

class OnlineASR:

def __init__(self, context_length=ASR_CONTEXT_LENGTH, language='en', cached=False):

"""

context_length: number of words to condition on

"""

self.context_length = context_length

self.audio_buffer = None

self.h_buffer = None

self.asr = OfflineASR(language, cached=cached)

self.reset()

def reset(self):

self.audio_buffer = AudioBuffer()

self.h_buffer = HypothesisBuffer()

@logger.catch

def process_chunk(self, chunk, finalize=False, return_audio=False):

sr = self.sample_rate

if finalize:

audio, buffer_offset = self.audio_buffer.clear()

logger.debug("Flushing audio buffer of length: {:.2f}", len(audio) / sr)

if not audio.any(): # buffer is empty

return None

else:

audio, buffer_offset = self.audio_buffer.push(chunk)

if audio is None: # buffer is not filled yet

return None

buffer_duration = len(audio) / sr

buffer_end_time = buffer_offset + buffer_duration

logger.debug("Transcribing audio of length: {:.2f}, buffer offset: {:.2f}", buffer_duration, buffer_offset)

if self.context_length > 0:

context = Word.to_text(self.h_buffer.confirmed_words)[-self.context_length:]

logger.debug("Conditioning on: {}", context)

else:

context = None

words = self.asr.transcribe(audio, previous_text=context)

words = Word.apply_offset(words, buffer_offset)

logger.opt(colors=True).debug("<g>Buffer transcription: {}</g>", Word.to_text(words))

# TODO: it relays on the last word end time, which is not accurate.

# better to track speech activity detection (vad)

if words:

silence_time = buffer_end_time - words[-1].end

else:

silence_time = buffer_duration

confirmed_words, unconfirmed_words = self.h_buffer.update(words)

if finalize:

self.h_buffer.clear()

if confirmed_words:

# fast forward to the middle of the last confirmed word and the first unconfirmed word

# it's less likely to cut the current utterance by mistake

start = confirmed_words[-1].end

end = unconfirmed_words[0].start if unconfirmed_words else buffer_end_time

ff_time = (start + end) / 2

self.audio_buffer.fast_forward(ff_time)

logger.debug("Fast forwarding audio buffer to: {:.2f}", ff_time)

result = {

"confirmed_text": Word.to_text(confirmed_words),

"unconfirmed_text": Word.to_text(unconfirmed_words),

"silence_time": silence_time

}

if return_audio:

result["audio"] = audio

logger.debug("Confirmed text: {}", result["confirmed_text"])

logger.debug("Unconfirmed text: {}", result["unconfirmed_text"])

logger.debug("Silence time: {:.2f}", silence_time)

return result

@property

def sample_rate(self):

return self.audio_buffer.sampling_rate

class ASRWithVAD:

def __init__(self, language='en', cached=False, vad_silence_threshold=VAD_SILENCE_STT_THRESHOLD):

self.asr = OfflineASR(language=language, cached=cached)

self.vad_silence_threshold = vad_silence_threshold

self.vad = VoiceActivityDetector()

self.audio_buffer = AudioBuffer(min_duration=1.0, max_duration=60)

def process_chunk(self, chunk, context=None):

vad_stats = self.vad.process_chunk(chunk)

audio_buffer, _ = self.audio_buffer.push(chunk)

if (audio_buffer is not None and

vad_stats["has_voice"] and

vad_stats["trailing_silence"] > self.vad_silence_threshold

):

logger.debug("STT silence threshold triggered, transcribing.")

buffer_text = self._transcribe(audio_buffer, context=context)

self.audio_buffer.clear()

self.vad.clear_voice()

else:

buffer_text = None

return {

"text": buffer_text,

"vad_stats": vad_stats

}

def _transcribe(self, audio_buffer, context=None):

if context is not None and context.last_message():

cond_text = context.last_message()["content"][0]["text"]

else:

cond_text = CONDITIONING_TEXT[self.asr.language]

words = self.asr.transcribe(audio_buffer, previous_text=cond_text)

return Word.to_text(words)

@property

def sample_rate(self):

return SAMPLING_RATE