makaveli10
382 lines
16 KiB
Python
382 lines
16 KiB
Python
import json
|
|
import logging
|
|
import threading
|
|
import time
|
|
import torch
|
|
|
|
from whisper_live.transcriber.transcriber_faster_whisper import WhisperModel
|
|
from whisper_live.backend.base import ServeClientBase
|
|
|
|
|
|
class ServeClientFasterWhisper(ServeClientBase):
|
|
|
|
SINGLE_MODEL = None
|
|
SINGLE_MODEL_LOCK = threading.Lock()
|
|
|
|
def __init__(self, websocket, task="transcribe", device=None, language=None, client_uid=None, model="small.en",
|
|
initial_prompt=None, vad_parameters=None, use_vad=True, single_model=False):
|
|
"""
|
|
Initialize a ServeClient instance.
|
|
The Whisper model is initialized based on the client's language and device availability.
|
|
The transcription thread is started upon initialization. A "SERVER_READY" message is sent
|
|
to the client to indicate that the server is ready.
|
|
|
|
Args:
|
|
websocket (WebSocket): The WebSocket connection for the client.
|
|
task (str, optional): The task type, e.g., "transcribe." Defaults to "transcribe".
|
|
device (str, optional): The device type for Whisper, "cuda" or "cpu". Defaults to None.
|
|
language (str, optional): The language for transcription. Defaults to None.
|
|
client_uid (str, optional): A unique identifier for the client. Defaults to None.
|
|
model (str, optional): The whisper model size. Defaults to 'small.en'
|
|
initial_prompt (str, optional): Prompt for whisper inference. Defaults to None.
|
|
single_model (bool, optional): Whether to instantiate a new model for each client connection. Defaults to False.
|
|
"""
|
|
super().__init__(client_uid, websocket)
|
|
self.model_sizes = [
|
|
"tiny", "tiny.en", "base", "base.en", "small", "small.en",
|
|
"medium", "medium.en", "large-v2", "large-v3", "distil-small.en",
|
|
"distil-medium.en", "distil-large-v2", "distil-large-v3",
|
|
"large-v3-turbo", "turbo"
|
|
]
|
|
|
|
self.model_size_or_path = model
|
|
self.language = "en" if self.model_size_or_path.endswith("en") else language
|
|
self.task = task
|
|
self.initial_prompt = initial_prompt
|
|
self.vad_parameters = vad_parameters or {"onset": 0.5}
|
|
self.no_speech_thresh = 0.45
|
|
self.same_output_threshold = 10
|
|
self.end_time_for_same_output = None
|
|
|
|
device = "cuda" if torch.cuda.is_available() else "cpu"
|
|
if device == "cuda":
|
|
major, _ = torch.cuda.get_device_capability(device)
|
|
self.compute_type = "float16" if major >= 7 else "float32"
|
|
else:
|
|
self.compute_type = "int8"
|
|
|
|
if self.model_size_or_path is None:
|
|
return
|
|
logging.info(f"Using Device={device} with precision {self.compute_type}")
|
|
|
|
try:
|
|
if single_model:
|
|
if ServeClientFasterWhisper.SINGLE_MODEL is None:
|
|
self.create_model(device)
|
|
ServeClientFasterWhisper.SINGLE_MODEL = self.transcriber
|
|
else:
|
|
self.transcriber = ServeClientFasterWhisper.SINGLE_MODEL
|
|
else:
|
|
self.create_model(device)
|
|
except Exception as e:
|
|
logging.error(f"Failed to load model: {e}")
|
|
self.websocket.send(json.dumps({
|
|
"uid": self.client_uid,
|
|
"status": "ERROR",
|
|
"message": f"Failed to load model: {str(self.model_size_or_path)}"
|
|
}))
|
|
self.websocket.close()
|
|
return
|
|
|
|
self.use_vad = use_vad
|
|
|
|
# threading
|
|
self.trans_thread = threading.Thread(target=self.speech_to_text)
|
|
self.trans_thread.start()
|
|
self.websocket.send(
|
|
json.dumps(
|
|
{
|
|
"uid": self.client_uid,
|
|
"message": self.SERVER_READY,
|
|
"backend": "faster_whisper"
|
|
}
|
|
)
|
|
)
|
|
|
|
def create_model(self, device):
|
|
"""
|
|
Instantiates a new model, sets it as the transcriber.
|
|
"""
|
|
self.transcriber = WhisperModel(
|
|
self.model_size_or_path,
|
|
device=device,
|
|
compute_type=self.compute_type,
|
|
local_files_only=False,
|
|
)
|
|
|
|
def check_valid_model(self, model_size):
|
|
"""
|
|
Check if it's a valid whisper model size.
|
|
|
|
Args:
|
|
model_size (str): The name of the model size to check.
|
|
|
|
Returns:
|
|
str: The model size if valid, None otherwise.
|
|
"""
|
|
if model_size not in self.model_sizes:
|
|
self.websocket.send(
|
|
json.dumps(
|
|
{
|
|
"uid": self.client_uid,
|
|
"status": "ERROR",
|
|
"message": f"Invalid model size {model_size}. Available choices: {self.model_sizes}"
|
|
}
|
|
)
|
|
)
|
|
return None
|
|
return model_size
|
|
|
|
def set_language(self, info):
|
|
"""
|
|
Updates the language attribute based on the detected language information.
|
|
|
|
Args:
|
|
info (object): An object containing the detected language and its probability. This object
|
|
must have at least two attributes: `language`, a string indicating the detected
|
|
language, and `language_probability`, a float representing the confidence level
|
|
of the language detection.
|
|
"""
|
|
if info.language_probability > 0.5:
|
|
self.language = info.language
|
|
logging.info(f"Detected language {self.language} with probability {info.language_probability}")
|
|
self.websocket.send(json.dumps(
|
|
{"uid": self.client_uid, "language": self.language, "language_prob": info.language_probability}))
|
|
|
|
def transcribe_audio(self, input_sample):
|
|
"""
|
|
Transcribes the provided audio sample using the configured transcriber instance.
|
|
|
|
If the language has not been set, it updates the session's language based on the transcription
|
|
information.
|
|
|
|
Args:
|
|
input_sample (np.array): The audio chunk to be transcribed. This should be a NumPy
|
|
array representing the audio data.
|
|
|
|
Returns:
|
|
The transcription result from the transcriber. The exact format of this result
|
|
depends on the implementation of the `transcriber.transcribe` method but typically
|
|
includes the transcribed text.
|
|
"""
|
|
if ServeClientFasterWhisper.SINGLE_MODEL:
|
|
ServeClientFasterWhisper.SINGLE_MODEL_LOCK.acquire()
|
|
result, info = self.transcriber.transcribe(
|
|
input_sample,
|
|
initial_prompt=self.initial_prompt,
|
|
language=self.language,
|
|
task=self.task,
|
|
vad_filter=self.use_vad,
|
|
vad_parameters=self.vad_parameters if self.use_vad else None)
|
|
if ServeClientFasterWhisper.SINGLE_MODEL:
|
|
ServeClientFasterWhisper.SINGLE_MODEL_LOCK.release()
|
|
|
|
if self.language is None and info is not None:
|
|
self.set_language(info)
|
|
return result
|
|
|
|
def get_previous_output(self):
|
|
"""
|
|
Retrieves previously generated transcription outputs if no new transcription is available
|
|
from the current audio chunks.
|
|
|
|
Checks the time since the last transcription output and, if it is within a specified
|
|
threshold, returns the most recent segments of transcribed text. It also manages
|
|
adding a pause (blank segment) to indicate a significant gap in speech based on a defined
|
|
threshold.
|
|
|
|
Returns:
|
|
segments (list): A list of transcription segments. This may include the most recent
|
|
transcribed text segments or a blank segment to indicate a pause
|
|
in speech.
|
|
"""
|
|
segments = []
|
|
if self.t_start is None:
|
|
self.t_start = time.time()
|
|
if time.time() - self.t_start < self.show_prev_out_thresh:
|
|
segments = self.prepare_segments()
|
|
|
|
# add a blank if there is no speech for 3 seconds
|
|
if len(self.text) and self.text[-1] != '':
|
|
if time.time() - self.t_start > self.add_pause_thresh:
|
|
self.text.append('')
|
|
return segments
|
|
|
|
def handle_transcription_output(self, result, duration):
|
|
"""
|
|
Handle the transcription output, updating the transcript and sending data to the client.
|
|
|
|
Args:
|
|
result (str): The result from whisper inference i.e. the list of segments.
|
|
duration (float): Duration of the transcribed audio chunk.
|
|
"""
|
|
segments = []
|
|
if len(result):
|
|
self.t_start = None
|
|
last_segment = self.update_segments(result, duration)
|
|
segments = self.prepare_segments(last_segment)
|
|
else:
|
|
# show previous output if there is pause i.e. no output from whisper
|
|
segments = self.get_previous_output()
|
|
|
|
if len(segments):
|
|
self.send_transcription_to_client(segments)
|
|
|
|
def speech_to_text(self):
|
|
"""
|
|
Process an audio stream in an infinite loop, continuously transcribing the speech.
|
|
|
|
This method continuously receives audio frames, performs real-time transcription, and sends
|
|
transcribed segments to the client via a WebSocket connection.
|
|
|
|
If the client's language is not detected, it waits for 30 seconds of audio input to make a language prediction.
|
|
It utilizes the Whisper ASR model to transcribe the audio, continuously processing and streaming results. Segments
|
|
are sent to the client in real-time, and a history of segments is maintained to provide context.Pauses in speech
|
|
(no output from Whisper) are handled by showing the previous output for a set duration. A blank segment is added if
|
|
there is no speech for a specified duration to indicate a pause.
|
|
|
|
Raises:
|
|
Exception: If there is an issue with audio processing or WebSocket communication.
|
|
|
|
"""
|
|
while True:
|
|
if self.exit:
|
|
logging.info("Exiting speech to text thread")
|
|
break
|
|
|
|
if self.frames_np is None:
|
|
continue
|
|
|
|
self.clip_audio_if_no_valid_segment()
|
|
|
|
input_bytes, duration = self.get_audio_chunk_for_processing()
|
|
if duration < 1.0:
|
|
time.sleep(0.1) # wait for audio chunks to arrive
|
|
continue
|
|
try:
|
|
input_sample = input_bytes.copy()
|
|
result = self.transcribe_audio(input_sample)
|
|
|
|
if result is None or self.language is None:
|
|
self.timestamp_offset += duration
|
|
time.sleep(0.25) # wait for voice activity, result is None when no voice activity
|
|
continue
|
|
self.handle_transcription_output(result, duration)
|
|
|
|
except Exception as e:
|
|
logging.error(f"[ERROR]: Failed to transcribe audio chunk: {e}")
|
|
time.sleep(0.01)
|
|
|
|
def format_segment(self, start, end, text, completed=False):
|
|
"""
|
|
Formats a transcription segment with precise start and end times alongside the transcribed text.
|
|
|
|
Args:
|
|
start (float): The start time of the transcription segment in seconds.
|
|
end (float): The end time of the transcription segment in seconds.
|
|
text (str): The transcribed text corresponding to the segment.
|
|
|
|
Returns:
|
|
dict: A dictionary representing the formatted transcription segment, including
|
|
'start' and 'end' times as strings with three decimal places and the 'text'
|
|
of the transcription.
|
|
"""
|
|
return {
|
|
'start': "{:.3f}".format(start),
|
|
'end': "{:.3f}".format(end),
|
|
'text': text,
|
|
'completed': completed
|
|
}
|
|
|
|
def update_segments(self, segments, duration):
|
|
"""
|
|
Processes the segments from whisper. Appends all the segments to the list
|
|
except for the last segment assuming that it is incomplete.
|
|
|
|
Updates the ongoing transcript with transcribed segments, including their start and end times.
|
|
Complete segments are appended to the transcript in chronological order. Incomplete segments
|
|
(assumed to be the last one) are processed to identify repeated content. If the same incomplete
|
|
segment is seen multiple times, it updates the offset and appends the segment to the transcript.
|
|
A threshold is used to detect repeated content and ensure it is only included once in the transcript.
|
|
The timestamp offset is updated based on the duration of processed segments. The method returns the
|
|
last processed segment, allowing it to be sent to the client for real-time updates.
|
|
|
|
Args:
|
|
segments(dict) : dictionary of segments as returned by whisper
|
|
duration(float): duration of the current chunk
|
|
|
|
Returns:
|
|
dict or None: The last processed segment with its start time, end time, and transcribed text.
|
|
Returns None if there are no valid segments to process.
|
|
"""
|
|
offset = None
|
|
self.current_out = ''
|
|
last_segment = None
|
|
|
|
# process complete segments
|
|
if len(segments) > 1 and segments[-1].no_speech_prob <= self.no_speech_thresh:
|
|
for i, s in enumerate(segments[:-1]):
|
|
text_ = s.text
|
|
self.text.append(text_)
|
|
with self.lock:
|
|
start, end = self.timestamp_offset + s.start, self.timestamp_offset + min(duration, s.end)
|
|
|
|
if start >= end:
|
|
continue
|
|
if s.no_speech_prob > self.no_speech_thresh:
|
|
continue
|
|
|
|
self.transcript.append(self.format_segment(start, end, text_, completed=True))
|
|
offset = min(duration, s.end)
|
|
|
|
# only process the last segment if it satisfies the no_speech_thresh
|
|
if segments[-1].no_speech_prob <= self.no_speech_thresh:
|
|
self.current_out += segments[-1].text
|
|
with self.lock:
|
|
last_segment = self.format_segment(
|
|
self.timestamp_offset + segments[-1].start,
|
|
self.timestamp_offset + min(duration, segments[-1].end),
|
|
self.current_out,
|
|
completed=False
|
|
)
|
|
|
|
if self.current_out.strip() == self.prev_out.strip() and self.current_out != '':
|
|
self.same_output_count += 1
|
|
|
|
# if we remove the audio because of same output on the nth reptition we might remove the
|
|
# audio thats not yet transcribed so, capturing the time when it was repeated for the first time
|
|
if self.end_time_for_same_output is None:
|
|
self.end_time_for_same_output = segments[-1].end
|
|
time.sleep(0.1) # wait for some voice activity just in case there is an unitended pause from the speaker for better punctuations.
|
|
else:
|
|
self.same_output_count = 0
|
|
self.end_time_for_same_output = None
|
|
|
|
# if same incomplete segment is seen multiple times then update the offset
|
|
# and append the segment to the list
|
|
if self.same_output_count > self.same_output_threshold:
|
|
if not len(self.text) or self.text[-1].strip().lower() != self.current_out.strip().lower():
|
|
self.text.append(self.current_out)
|
|
with self.lock:
|
|
self.transcript.append(self.format_segment(
|
|
self.timestamp_offset,
|
|
self.timestamp_offset + min(duration, self.end_time_for_same_output),
|
|
self.current_out,
|
|
completed=True
|
|
))
|
|
self.current_out = ''
|
|
offset = min(duration, self.end_time_for_same_output)
|
|
self.same_output_count = 0
|
|
last_segment = None
|
|
self.end_time_for_same_output = None
|
|
else:
|
|
self.prev_out = self.current_out
|
|
|
|
# update offset
|
|
if offset is not None:
|
|
with self.lock:
|
|
self.timestamp_offset += offset
|
|
|
|
return last_segment
|
|
|