feat: add advanced tuning, diarization, and LLM support to multi-process engine
This commit is contained in:
Adolfo Reyna
+79
-44
@@ -16,13 +16,14 @@ except ImportError:
|
||||
sys.modules["_lzma"] = MagicMock()
|
||||
sys.modules["lzma"] = mock_lzma
|
||||
|
||||
import sys
|
||||
import time
|
||||
import os
|
||||
import numpy as np
|
||||
import sounddevice as sd
|
||||
import torch
|
||||
import mlx_whisper
|
||||
from silero_vad import load_silero_vad, get_speech_timestamps
|
||||
from pyannote.audio import Pipeline
|
||||
import queue
|
||||
import argparse
|
||||
|
||||
@@ -36,7 +37,7 @@ def transcribe_with_controls(audio, transcribe_kwargs):
|
||||
try:
|
||||
return mlx_whisper.transcribe(audio, **transcribe_kwargs)
|
||||
except Exception as exc:
|
||||
# Fallback logic for beam_size or other specific MLX implementation gaps
|
||||
message = str(exc)
|
||||
if "beam_size" in transcribe_kwargs:
|
||||
transcribe_kwargs.pop("beam_size")
|
||||
return mlx_whisper.transcribe(audio, **transcribe_kwargs)
|
||||
@@ -47,14 +48,44 @@ def audio_callback(indata, frames, time, status, audio_queue):
|
||||
print(status, file=sys.stderr)
|
||||
audio_queue.put(indata.copy())
|
||||
|
||||
def assign_speakers_to_segments(segments, diarization):
|
||||
if not diarization or not segments:
|
||||
return segments
|
||||
for segment in segments:
|
||||
segment_start, segment_end = segment['start'], segment['end']
|
||||
speaker_intersections = {}
|
||||
for turn, _, speaker in diarization.itertracks(yield_label=True):
|
||||
intersection_start = max(segment_start, turn.start)
|
||||
intersection_end = min(segment_end, turn.end)
|
||||
if intersection_end > intersection_start:
|
||||
duration = intersection_end - intersection_start
|
||||
speaker_intersections[speaker] = speaker_intersections.get(speaker, 0) + duration
|
||||
if speaker_intersections:
|
||||
segment['speaker'] = max(speaker_intersections, key=speaker_intersections.get)
|
||||
else:
|
||||
segment['speaker'] = 'UNKNOWN'
|
||||
return segments
|
||||
|
||||
def run_transcription(out_queue, args):
|
||||
device = "mps" if torch.backends.mps.is_available() else "cpu"
|
||||
print(f"[Transcribe] Loading Whisper model '{args.model}' on {device}...")
|
||||
model_name = args.model
|
||||
if args.quantize and "4bit" not in model_name:
|
||||
model_name = "mlx-community/whisper-small-mlx-4bit"
|
||||
|
||||
print(f"[Transcribe] Loading Whisper model '{model_name}' on {device}...")
|
||||
vad_model = load_silero_vad()
|
||||
audio_queue = queue.Queue()
|
||||
|
||||
# Selection of device
|
||||
diarization_pipeline = None
|
||||
if args.speaker_diarization:
|
||||
print("[Transcribe] Loading speaker diarization pipeline...")
|
||||
hf_token = os.environ.get("HF_TOKEN")
|
||||
try:
|
||||
diarization_pipeline = Pipeline.from_pretrained("pyannote/speaker-diarization-3.1", use_auth_token=hf_token)
|
||||
diarization_pipeline.to(torch.device(device))
|
||||
except Exception as e:
|
||||
print(f"[Transcribe] Diarization load failed: {e}")
|
||||
|
||||
audio_queue = queue.Queue()
|
||||
device_index = args.device
|
||||
|
||||
def callback_wrapper(indata, frames, time, status):
|
||||
@@ -63,16 +94,17 @@ def run_transcription(out_queue, args):
|
||||
audio_buffer = []
|
||||
speech_started = False
|
||||
buffer_limit = SAMPLERATE * args.max_buffer
|
||||
rolling_context = ""
|
||||
last_stream_time = time.time()
|
||||
|
||||
print(f"[Transcribe] Starting audio stream on device {device_index}...")
|
||||
print(f"[Transcribe] Audio stream ready on device {device_index}.")
|
||||
|
||||
try:
|
||||
with sd.InputStream(samplerate=SAMPLERATE, channels=args.channels, callback=callback_wrapper, blocksize=BLOCK_SIZE, device=device_index):
|
||||
while True:
|
||||
while not audio_queue.empty():
|
||||
data = audio_queue.get()
|
||||
if args.channels > 1:
|
||||
data = np.mean(data, axis=1) # Mix to mono
|
||||
if args.channels > 1: data = np.mean(data, axis=1)
|
||||
audio_buffer.append(data.flatten())
|
||||
|
||||
if audio_buffer:
|
||||
@@ -80,64 +112,67 @@ def run_transcription(out_queue, args):
|
||||
audio_tensor = torch.from_numpy(current_audio)
|
||||
buffer_duration = len(current_audio) / SAMPLERATE
|
||||
|
||||
speech_timestamps = get_speech_timestamps(
|
||||
audio_tensor,
|
||||
vad_model,
|
||||
sampling_rate=SAMPLERATE,
|
||||
threshold=VAD_THRESHOLD,
|
||||
min_silence_duration_ms=args.silence
|
||||
)
|
||||
speech_timestamps = get_speech_timestamps(audio_tensor, vad_model, sampling_rate=SAMPLERATE, threshold=VAD_THRESHOLD, min_silence_duration_ms=args.silence)
|
||||
|
||||
if speech_timestamps:
|
||||
speech_started = True
|
||||
last_end = speech_timestamps[-1]['end']
|
||||
buffer_len_samples = len(current_audio)
|
||||
|
||||
# Check for flush
|
||||
if (buffer_len_samples - last_end) > (SAMPLERATE * args.silence / 1000) or buffer_len_samples > buffer_limit:
|
||||
# 1. Transcribe
|
||||
transcribe_kwargs = {
|
||||
"path_or_hf_repo": args.model,
|
||||
"temperature": (0.0, 0.2, 0.4, 0.6, 0.8, 1.0),
|
||||
"path_or_hf_repo": model_name,
|
||||
"temperature": args.temperature_fallback,
|
||||
"word_timestamps": args.speaker_diarization,
|
||||
"logprob_threshold": args.logprob_threshold,
|
||||
"compression_ratio_threshold": args.compression_threshold,
|
||||
}
|
||||
if args.lang:
|
||||
transcribe_kwargs["language"] = args.lang
|
||||
if args.lang: transcribe_kwargs["language"] = args.lang
|
||||
if args.context and rolling_context: transcribe_kwargs["initial_prompt"] = rolling_context
|
||||
|
||||
result = transcribe_with_controls(current_audio, transcribe_kwargs)
|
||||
text = result['text'].strip()
|
||||
detected_lang = result.get('language', 'en')
|
||||
detected_lang = result.get('language', args.lang or 'en')
|
||||
|
||||
if args.filter_lang and args.lang and detected_lang != args.lang:
|
||||
print(f"[Transcribe] Filtered {detected_lang}")
|
||||
text = ""
|
||||
|
||||
if text:
|
||||
# Send to Translation Process
|
||||
out_queue.put({
|
||||
"original": text,
|
||||
"detected_lang": detected_lang,
|
||||
"ts": time.time()
|
||||
})
|
||||
segments = result.get('segments', [])
|
||||
if diarization_pipeline:
|
||||
audio_for_diarization = torch.from_numpy(current_audio).float().unsqueeze(0)
|
||||
try:
|
||||
diarization_result = diarization_pipeline({"waveform": audio_for_diarization, "sample_rate": SAMPLERATE})
|
||||
segments = assign_speakers_to_speakers(segments, diarization_result)
|
||||
except: pass
|
||||
|
||||
out_queue.put({"original": text, "detected_lang": detected_lang, "segments": segments, "ts": time.time()})
|
||||
print(f"[Transcribe] {detected_lang.upper()}: {text}")
|
||||
if args.context: rolling_context = (rolling_context + " " + text)[-200:].strip()
|
||||
|
||||
audio_buffer = []
|
||||
speech_started = False
|
||||
|
||||
elif args.stream and (time.time() - last_stream_time) > 1.5:
|
||||
# Draft Mode
|
||||
draft_kwargs = {"path_or_hf_repo": model_name, "temperature": 0.0}
|
||||
if args.lang: draft_kwargs["language"] = args.lang
|
||||
if args.context and rolling_context: draft_kwargs["initial_prompt"] = rolling_context
|
||||
|
||||
draft_result = transcribe_with_controls(current_audio, draft_kwargs)
|
||||
draft_text = draft_result['text'].strip()
|
||||
if draft_text:
|
||||
out_queue.put({"draft": draft_text, "ts": time.time()})
|
||||
last_stream_time = time.time()
|
||||
|
||||
elif not speech_started and len(current_audio) > SAMPLERATE * 2:
|
||||
audio_buffer = []
|
||||
|
||||
except KeyboardInterrupt:
|
||||
print("[Transcribe] Stopped.")
|
||||
except Exception as e:
|
||||
print(f"[Transcribe] Error: {e}")
|
||||
except KeyboardInterrupt: pass
|
||||
except Exception as e: print(f"[Transcribe] Error: {e}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--model", type=str, default="mlx-community/whisper-base-mlx")
|
||||
parser.add_argument("--device", type=int, default=None)
|
||||
parser.add_argument("--lang", type=str, default=None)
|
||||
parser.add_argument("--silence", type=int, default=1000)
|
||||
parser.add_argument("--max-buffer", type=int, default=20)
|
||||
parser.add_argument("--channels", type=int, default=1)
|
||||
args = parser.parse_args()
|
||||
|
||||
# This part would normally be called by the main coordinator
|
||||
# For testing, we can use a dummy queue
|
||||
import multiprocessing
|
||||
q = multiprocessing.Queue()
|
||||
run_transcription(q, args)
|
||||
run_transcription(multiprocessing.Queue(), argparse.Namespace())
|
||||
|
||||
Reference in New Issue
Block a user