adolforeyna/whisper-translation
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: add advanced tuning, diarization, and LLM support to multi-process engine

Browse Source
This commit is contained in:
Adolfo Reyna
2026-03-15 19:53:36 -04:00
parent c474184169
commit ca07d96ff0
5 changed files with 188 additions and 113 deletions
Download Patch File Download Diff File Expand all files Collapse all files
config.json
+17 -1
View File
@@ -8,5 +8,21 @@
"es": false,
"fr": false,
"ar": false,
"ingest": false
"en": false,
"ingest": false,
"filter_lang": false,
"stream": false,
"context": false,
"quantize": false,
"speaker_diarization": false,
"post_correct": false,
"post_correct_llm": false,
"post_correct_model": "qwen:2b",
"llm_paragraph": false,
"temperature_fallback": [0.0, 0.2, 0.4, 0.6, 0.8, 1.0],
"logprob_threshold": -0.8,
"compression_threshold": 2.2,
"mt_max_chars": 250,
"mt_max_new_tokens": 150,
"mt_num_beams": 4
}
engine_distribute.py
+21 -19
View File
@@ -2,6 +2,7 @@ import requests
import time
import argparse
import multiprocessing
import sys
INGEST_URL = "https://emiapi.reynafamily.com/live-captions/ingest"
@@ -13,37 +14,38 @@ def run_distribution(in_queue, args):
payload = in_queue.get()
if payload is None: break
# Draft support
if "draft" in payload:
draft_text = payload["draft"]
sys.stdout.write(f"\r\033[K[DRAFT]: {draft_text}")
sys.stdout.flush()
if args.ingest:
delay = 1
max_delay = 15
success = False
requests.post(INGEST_URL, json={"draft": draft_text}, timeout=2)
continue
# Finalized segment
print(f"\n[Final]: {payload.get('original', '')}")
for lang in ["es", "fr", "ar", "en"]:
if payload.get(lang):
print(f"[{lang.upper()}]: {payload[lang]}")
if args.ingest:
delay, max_delay, success = 1, 15, False
while not success:
try:
response = requests.post(INGEST_URL, json=payload, timeout=5)
if response.status_code == 200:
success = True
else:
print(f"[Distribute Error] {response.status_code}. Retrying in {delay}s...")
if response.status_code == 200: success = True
else: print(f"[Distribute Error] {response.status_code}. Retrying...")
except Exception as e:
print(f"[Distribute Error] {e}. Retrying in {delay}s...")
print(f"[Distribute Error] {e}. Retrying...")
if not success:
time.sleep(delay)
delay = min(delay * 2, max_delay)
print(f"[Distribute] Sent payload: {list(payload.keys())}")
else:
print(f"[Distribute] Local display (ingest disabled): {payload}")
except Exception as e:
print(f"[Distribute] Error: {e}")
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("-i", "--ingest", action="store_true")
args = parser.parse_args()
# Dummy queue for testing
in_q = multiprocessing.Queue()
run_distribution(in_q, args)
import multiprocessing
run_distribution(multiprocessing.Queue(), argparse.Namespace())
engine_transcribe.py
+79 -44
View File
@@ -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())
engine_translate.py
+43 -40
View File
@@ -22,6 +22,7 @@ import re
import argparse
import multiprocessing
import time
import requests
TARGET_LANGS = {
"es": "Helsinki-NLP/opus-mt-en-es",
@@ -33,98 +34,100 @@ def split_text_for_translation(text, max_chars=250):
normalized = " ".join(text.split()).strip()
if not normalized or len(normalized) <= max_chars:
return [normalized] if normalized else []
chunks = []
current = ""
sentences = [s for s in re.split(r"(?<=[.!?])\s+", normalized) if s]
for sentence in sentences:
if len(sentence) > max_chars:
words = sentence.split()
word_chunk = ""
for word in words:
candidate = f"{word_chunk} {word}".strip()
if len(candidate) <= max_chars:
word_chunk = candidate
if len(candidate) <= max_chars: word_chunk = candidate
else:
if word_chunk: chunks.append(word_chunk)
word_chunk = word
if word_chunk: chunks.append(word_chunk)
continue
candidate = f"{current} {sentence}".strip()
if candidate and len(candidate) <= max_chars:
current = candidate
if candidate and len(candidate) <= max_chars: current = candidate
else:
if current: chunks.append(current)
current = sentence
if current: chunks.append(current)
return chunks
def call_ollama_generate(args, prompt):
payload = {"model": args.post_correct_model, "prompt": prompt, "stream": False}
try:
resp = requests.post("http://127.0.0.1:11434/api/generate", json=payload, timeout=10)
return resp.json().get("response", "").strip()
except: return ""
def structure_paragraph_with_llm(text, args):
prompt = f"Format the following transcribed segments into a coherent paragraph. Output ONLY the paragraph:\n{text}"
return call_ollama_generate(args, prompt) or text
def post_correct_with_llm(text, args):
prompt = f"Correct grammar and typos in this caption. Output ONLY the corrected text:\n{text}"
return call_ollama_generate(args, prompt) or text
def run_translation(in_queue, out_queue, args):
device = "mps" if torch.backends.mps.is_available() else "cpu"
print(f"[Translate] Loading translation models on {device}...")
translation_engines = {}
for lang_key, model_id in TARGET_LANGS.items():
if getattr(args, lang_key, False):
print(f"[Translate] Loading {lang_key} model...")
tokenizer = MarianTokenizer.from_pretrained(model_id)
model = MarianMTModel.from_pretrained(model_id).to(device)
translation_engines[lang_key] = (model, tokenizer)
translation_engines[lang_key] = (
MarianMTModel.from_pretrained(model_id).to(device),
MarianTokenizer.from_pretrained(model_id)
)
paragraph_buffer = []
print("[Translate] Ready.")
while True:
try:
item = in_queue.get()
if item is None: break
# Draft support
if "draft" in item:
out_queue.put({"draft": item["draft"]})
continue
text = item.get("original", "").strip()
detected_lang = item.get("detected_lang", "en")
# Simple paragraph logic:
# If the segment ends with sentence-terminal punctuation, flush the paragraph.
if args.post_correct and args.post_correct_llm:
text = post_correct_with_llm(text, args)
paragraph_buffer.append(text)
if text.endswith((".", "?", "!")):
# Paragraph logic
if not args.llm_paragraph or text.endswith((".", "?", "!")):
full_text = " ".join(paragraph_buffer)
payload = {"original": full_text, "en": full_text if detected_lang == "en" else None}
# If detected language is not English, we'd normally bridge to English first.
# For now, let's assume direct translation for simplicity or bridge if needed.
# (Refining bridge logic can come later)
if args.llm_paragraph:
full_text = structure_paragraph_with_llm(full_text, args)
payload = {"original": full_text, "ts": item.get("ts", time.time())}
if args.en: payload["en"] = full_text
for lang_key, (model, tokenizer) in translation_engines.items():
chunks = split_text_for_translation(full_text)
chunks = split_text_for_translation(full_text, args.mt_max_chars)
translated_parts = []
for chunk in chunks:
inputs = tokenizer(chunk, return_tensors="pt", padding=True).to(device)
with torch.no_grad():
translated_tokens = model.generate(**inputs, max_new_tokens=150)
translated_tokens = model.generate(**inputs, max_new_tokens=args.mt_max_new_tokens, num_beams=args.mt_num_beams)
translated_parts.append(tokenizer.decode(translated_tokens[0], skip_special_tokens=True).strip())
translated_text = " ".join(part for part in translated_parts if part).strip()
payload[lang_key] = translated_text
print(f"[Translate] {lang_key.upper()}: {translated_text}")
payload[lang_key] = " ".join(translated_parts)
out_queue.put(payload)
paragraph_buffer = []
except Exception as e:
print(f"[Translate] Error: {e}")
except Exception as e: print(f"[Translate] Error: {e}")
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("-es", action="store_true")
parser.add_argument("-fr", action="store_true")
parser.add_argument("-ar", action="store_true")
args = parser.parse_args()
# Dummy queues for testing
in_q = multiprocessing.Queue()
out_q = multiprocessing.Queue()
run_translation(in_q, out_q, args)
import multiprocessing
run_translation(multiprocessing.Queue(), multiprocessing.Queue(), argparse.Namespace())
main_v2.py
+27 -8
View File
@@ -18,13 +18,18 @@ except ImportError:
import multiprocessing
import argparse
import sys
import json
import os
from engine_transcribe import run_transcription
from engine_translate import run_translation
from engine_distribute import run_distribution
def parse_temperature_fallback(value):
try:
return tuple(float(x.strip()) for x in value.split(",") if x.strip())
except:
return (0.0, 0.2, 0.4, 0.6, 0.8, 1.0)
def main():
config_path = "config.json"
defaults = {}
@@ -41,21 +46,35 @@ def main():
parser.add_argument("--silence", type=int, default=defaults.get("silence", 1000))
parser.add_argument("--max-buffer", type=int, default=defaults.get("max_buffer", 20))
parser.add_argument("--channels", type=int, default=defaults.get("channels", 1))
parser.add_argument("--filter-lang", action="store_true", default=defaults.get("filter_lang", False))
parser.add_argument("-q", "--quantize", action="store_true", default=defaults.get("quantize", False))
parser.add_argument("-s", "--stream", action="store_true", default=defaults.get("stream", False))
parser.add_argument("-c", "--context", action="store_true", default=defaults.get("context", False))
parser.add_argument("--speaker-diarization", action="store_true", default=defaults.get("speaker_diarization", False))
parser.add_argument("--temperature-fallback", type=parse_temperature_fallback, default=tuple(defaults.get("temperature_fallback", [0.0, 0.2, 0.4, 0.6, 0.8, 1.0])))
parser.add_argument("--logprob-threshold", type=float, default=defaults.get("logprob_threshold", -0.8))
parser.add_argument("--compression-threshold", type=float, default=defaults.get("compression_threshold", 2.2))
# Translate Args
parser.add_argument("-es", action="store_true", default=defaults.get("es", False), help="Enable Spanish translation")
parser.add_argument("-fr", action="store_true", default=defaults.get("fr", False), help="Enable French translation")
parser.add_argument("-ar", action="store_true", default=defaults.get("ar", False), help="Enable Arabic translation")
parser.add_argument("-es", action="store_true", default=defaults.get("es", False))
parser.add_argument("-fr", action="store_true", default=defaults.get("fr", False))
parser.add_argument("-ar", action="store_true", default=defaults.get("ar", False))
parser.add_argument("-en", action="store_true", default=defaults.get("en", False))
parser.add_argument("--post-correct", action="store_true", default=defaults.get("post_correct", False))
parser.add_argument("--post-correct-llm", action="store_true", default=defaults.get("post_correct_llm", False))
parser.add_argument("--post-correct-model", type=str, default=defaults.get("post_correct_model", "qwen:2b"))
parser.add_argument("--llm-paragraph", action="store_true", default=defaults.get("llm_paragraph", False))
parser.add_argument("--mt-max-chars", type=int, default=defaults.get("mt_max_chars", 250))
parser.add_argument("--mt-max-new-tokens", type=int, default=defaults.get("mt_max_new_tokens", 150))
parser.add_argument("--mt-num-beams", type=int, default=defaults.get("mt_num_beams", 4))
# Distribute Args
parser.add_argument("-i", "--ingest", action="store_true", default=defaults.get("ingest", False), help="Enable data transmission to server")
parser.add_argument("-i", "--ingest", action="store_true", default=defaults.get("ingest", False))
args = parser.parse_args()
# Queues for communication
# Transcribe -> Translate
# Queues
q_trans_to_tl = multiprocessing.Queue()
# Translate -> Distribute
q_tl_to_dist = multiprocessing.Queue()
# Processes