Files
whisper-translation/transcribe.py
T
2026-03-09 21:33:11 -04:00

1085 lines
54 KiB
Python

import sys
import time
import re
import os
import shutil
from datetime import datetime
import requests
import threading
import json
import argparse
from unittest.mock import MagicMock
from collections import Counter, deque
# Comprehensive workaround for missing _lzma in some Python builds
try:
import lzma
except ImportError:
mock_lzma = MagicMock()
mock_lzma.FORMAT_XZ = 1
mock_lzma.FORMAT_ALONE = 2
mock_lzma.FORMAT_RAW = 3
mock_lzma.CHECK_NONE = 0
mock_lzma.CHECK_CRC32 = 1
mock_lzma.CHECK_CRC64 = 4
mock_lzma.CHECK_SHA256 = 10
sys.modules["_lzma"] = MagicMock()
sys.modules["lzma"] = mock_lzma
import mlx_whisper
import numpy as np
import sounddevice as sd
import queue
import torch
from silero_vad import load_silero_vad, get_speech_timestamps
from transformers import MarianMTModel, MarianTokenizer
# Parameters
WHISPER_MODEL = "mlx-community/whisper-base-mlx"
INGEST_URL = "https://emiapi.reynafamily.com/live-captions/ingest"
# Translation models (English -> Target)
# Map to the specific keys requested by the backend
TARGET_LANGS = {
"es": "Helsinki-NLP/opus-mt-en-es",
"fr": "Helsinki-NLP/opus-mt-en-fr",
"ar": "Helsinki-NLP/opus-mt-en-ar" # Added Arabic as discussed before
}
SAMPLERATE = 16000
BLOCK_SIZE = 512
VAD_THRESHOLD = 0.5
audio_queue = queue.Queue()
ingest_queue = queue.Queue()
session_log_handle = None
def log_session(message):
"""Write a timestamped line to the session debug log if enabled."""
global session_log_handle
if session_log_handle is None:
return
ts = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
session_log_handle.write(f"{ts} {message}\n")
session_log_handle.flush()
def parse_temperature_fallback(value):
"""Parse comma-separated temperatures into a tuple of floats."""
try:
temps = tuple(float(x.strip()) for x in value.split(",") if x.strip())
except ValueError as exc:
raise argparse.ArgumentTypeError("Invalid --temperature-fallback value.") from exc
if not temps:
raise argparse.ArgumentTypeError("--temperature-fallback requires at least one value.")
return temps
def split_text_for_translation(text, max_chars=250):
"""Split long English text into sentence-aware chunks for MT."""
normalized = " ".join(text.split()).strip()
if not normalized:
return []
if len(normalized) <= max_chars:
return [normalized]
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
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
else:
if current:
chunks.append(current)
current = sentence
if current:
chunks.append(current)
return chunks
def transcribe_with_controls(audio, transcribe_kwargs):
"""Call mlx_whisper.transcribe and gracefully fallback if a decoder arg is unsupported."""
optional_keys = [
"beam_size",
"temperature",
"logprob_threshold",
"compression_ratio_threshold",
]
try:
return mlx_whisper.transcribe(audio, **transcribe_kwargs)
except TypeError as exc:
message = str(exc)
unsupported = [k for k in optional_keys if f"'{k}'" in message]
if not unsupported:
raise
retry_kwargs = {k: v for k, v in transcribe_kwargs.items() if k not in unsupported}
print(f"\n[SYSTEM]: Decoder args not supported by current mlx_whisper build: {', '.join(unsupported)}. Retrying without them.")
return mlx_whisper.transcribe(audio, **retry_kwargs)
except Exception as exc:
message = str(exc).lower()
if "beam search decoder is not yet implemented" in message and "beam_size" in transcribe_kwargs:
retry_kwargs = dict(transcribe_kwargs)
retry_kwargs.pop("beam_size", None)
print("\n[SYSTEM]: Beam search is not implemented in this mlx_whisper build. Retrying with greedy decoding.")
return mlx_whisper.transcribe(audio, **retry_kwargs)
raise
def extract_whisper_quality(result):
"""Compute aggregate quality signals from Whisper segments when available."""
segments = result.get("segments") if isinstance(result, dict) else None
if not segments:
return None, None
logprobs = []
compressions = []
for seg in segments:
avg_logprob = seg.get("avg_logprob")
compression_ratio = seg.get("compression_ratio")
if isinstance(avg_logprob, (int, float)):
logprobs.append(float(avg_logprob))
if isinstance(compression_ratio, (int, float)):
compressions.append(float(compression_ratio))
mean_logprob = sum(logprobs) / len(logprobs) if logprobs else None
max_compression = max(compressions) if compressions else None
return mean_logprob, max_compression
def quality_score(mean_logprob, max_compression):
"""Higher score means better quality."""
if mean_logprob is None:
mean_logprob = -9.0
compression_penalty = 0.25 * max(0.0, (max_compression or 1.5) - 1.5)
return mean_logprob - compression_penalty
def cosine_similarity(vec_a, vec_b):
if vec_a is None or vec_b is None:
return None
denom = (np.linalg.norm(vec_a) * np.linalg.norm(vec_b))
if denom == 0:
return None
return float(np.dot(vec_a, vec_b) / denom)
def compute_speaker_signature(audio, samplerate):
"""Compute a lightweight speaker signature from averaged log spectrum."""
if audio is None or len(audio) < int(0.5 * samplerate):
return None
data = np.asarray(audio, dtype=np.float32)
if data.ndim > 1:
data = np.mean(data, axis=1)
# Keep only voiced-ish samples to reduce silence/noise influence.
voiced = data[np.abs(data) > 0.01]
if len(voiced) < int(0.3 * samplerate):
voiced = data
frame_size = 1024
hop = 512
if len(voiced) < frame_size:
return None
window = np.hanning(frame_size).astype(np.float32)
spectra = []
for i in range(0, len(voiced) - frame_size + 1, hop):
frame = voiced[i:i + frame_size] * window
mag = np.abs(np.fft.rfft(frame))
if mag.size > 1:
mag = mag[1:257] # Focus on lower-frequency envelope.
spectra.append(np.log1p(mag))
if not spectra:
return None
signature = np.mean(np.stack(spectra, axis=0), axis=0)
norm = np.linalg.norm(signature)
if norm == 0:
return None
return signature / norm
def should_retry_transcription(text, mean_logprob, max_compression, args):
if not text:
return False
if mean_logprob is not None and mean_logprob < args.retry_logprob_threshold:
return True
if max_compression is not None and max_compression > args.retry_compression_threshold:
return True
return False
def parse_glossary_pair(value):
if "=" not in value:
raise argparse.ArgumentTypeError("Glossary pairs must be in SOURCE=TARGET format.")
source, target = value.split("=", 1)
source = source.strip()
target = target.strip()
if not source or not target:
raise argparse.ArgumentTypeError("Glossary SOURCE and TARGET cannot be empty.")
return source, target
def apply_glossary(text, glossary_pairs):
updated = text
for source, target in glossary_pairs:
pattern = re.compile(rf"\b{re.escape(source)}\b", re.IGNORECASE)
updated = pattern.sub(target, updated)
return updated
def load_glossary_file(path):
"""Load SOURCE=TARGET glossary pairs from a file."""
pairs = []
if not path or not os.path.exists(path):
return pairs
with open(path, "r", encoding="utf-8") as f:
for line_no, raw in enumerate(f, start=1):
line = raw.strip()
if not line or line.startswith("#"):
continue
try:
pairs.append(parse_glossary_pair(line))
except argparse.ArgumentTypeError:
print(f"[SYSTEM]: Skipping invalid glossary line {line_no} in {path}: {line}")
return pairs
def normalize_english_caption(text):
normalized = " ".join(text.split()).strip()
if not normalized:
return ""
if normalized[0].isalpha():
normalized = normalized[0].upper() + normalized[1:]
return normalized
def apply_rule_based_post_correction(text):
"""Apply lightweight deterministic cleanup for finalized English captions."""
if not text:
return ""
corrected = normalize_english_caption(text)
corrected = re.sub(r"\s+", " ", corrected).strip()
corrected = re.sub(r"\s+([,.;!?])", r"\1", corrected)
corrected = re.sub(r"([,.;!?]){2,}", r"\1", corrected)
# Remove common disfluencies conservatively.
corrected = re.sub(r"\b(uh+|um+|erm+|ah+|hmm+)\b", "", corrected, flags=re.IGNORECASE)
corrected = re.sub(r"\s+", " ", corrected).strip()
# Collapse immediate duplicated words (e.g., "the the", "I I").
corrected = re.sub(r"\b(\w+)\s+\1\b", r"\1", corrected, flags=re.IGNORECASE)
corrected = re.sub(r"\s+", " ", corrected).strip()
if corrected and corrected[0].isalpha():
corrected = corrected[0].upper() + corrected[1:]
return corrected
def word_overlap_ratio(source_text, candidate_text):
"""Compute token overlap ratio between source and candidate."""
src_tokens = set(re.findall(r"[a-z0-9']+", source_text.lower()))
cand_tokens = set(re.findall(r"[a-z0-9']+", candidate_text.lower()))
if not src_tokens:
return 1.0
return len(src_tokens & cand_tokens) / len(src_tokens)
def call_ollama_generate(args, prompt, timeout, temperature=0.1):
"""Call local Ollama with keep_alive to reduce cold-start latency."""
payload = {
"model": args.post_correct_model,
"prompt": prompt,
"stream": False,
"options": {"temperature": temperature},
"keep_alive": args.post_correct_keep_alive,
}
return requests.post(
args.post_correct_ollama_url,
json=payload,
timeout=timeout,
)
def post_correct_with_local_llm(text, args, llm_state, context_lines=None):
"""Optionally rewrite finalized English caption via a local Ollama model."""
if args.post_correct_debug:
print(f"[POST-LLM][PRE]: {text}", flush=True)
if not args.post_correct_llm:
if args.post_correct_debug:
print("[POST-LLM][SKIP]: --post-correct-llm not enabled; using rules-only.", flush=True)
return text
if not llm_state.get("available", True):
if args.post_correct_debug:
print("[POST-LLM][SKIP]: LLM marked unavailable after prior failure; using rules-only.", flush=True)
return text
context_lines = context_lines or []
prev2 = context_lines[-2] if len(context_lines) >= 2 else ""
prev1 = context_lines[-1] if len(context_lines) >= 1 else ""
prompt = (
"You are a real-time English caption corrector for live speech.\n"
"Task:\n"
"Clean ONLY the current caption line.\n"
"Hard rules:\n"
"1. Preserve meaning exactly. Never replace current content with prior context.\n"
"2. Remove disfluencies and false starts (um, uh, you know, I mean, repeated starts).\n"
"3. Fix punctuation, casing, and obvious STT typos.\n"
"4. Keep the speaker's tone and intent; do not paraphrase for style.\n"
"5. If uncertain, return the original line unchanged.\n"
"6. Output ONLY one corrected line. No explanations, no quotes.\n"
"Context (reference only; never override current line):\n"
f"Previous line 1: {prev1}\n"
f"Previous line 2: {prev2}\n"
"Current line to correct:\n"
)
prompt += f"{text}\nCorrected:"
try:
response = call_ollama_generate(args, prompt, args.post_correct_llm_timeout, temperature=0.1)
response.raise_for_status()
raw = response.json().get("response", "").strip()
cleaned = normalize_english_caption(raw)
if args.post_correct_debug:
print(f"[POST-LLM][RAW]: {raw}", flush=True)
if cleaned:
overlap = word_overlap_ratio(text, cleaned)
if overlap < args.post_correct_min_overlap:
if args.post_correct_debug:
print(
f"[POST-LLM][REJECT]: low overlap ({overlap:.2f} < {args.post_correct_min_overlap:.2f}); keeping pre text"
, flush=True)
return text
if args.post_correct_debug:
print(f"[POST-LLM][POST]: {cleaned}", flush=True)
return cleaned
if args.post_correct_debug:
print("[POST-LLM][POST]: <empty response, keeping pre text>", flush=True)
except Exception as exc:
if not llm_state.get("warned", False):
print(f"\n[SYSTEM]: Local post-correct LLM unavailable ({exc}). Falling back to rules-only.")
llm_state["warned"] = True
llm_state["available"] = False
if args.post_correct_debug:
print(f"[POST-LLM][ERROR]: {exc}", flush=True)
return text
def maybe_merge_recent_caption(recent_en_lines, current_text, now_ts, window_sec):
if not recent_en_lines:
return None
previous = recent_en_lines[-1]
if (now_ts - previous["ts"]) > window_sec:
return None
prev_text = previous["text"].strip()
if not prev_text:
return None
sentence_end = bool(re.search(r'[.!?]["\')\]]*$', prev_text))
if sentence_end:
return None
if len(prev_text) > 120:
return None
merged = normalize_english_caption(f"{prev_text} {current_text}")
if merged == prev_text:
return None
old = previous["text"]
previous["text"] = merged
previous["ts"] = now_ts
return old, merged
def llm_decide_caption_merge(previous_text, current_text, candidate_text, args, llm_state):
"""Ask local LLM to approve/refine a merge candidate; fallback to heuristic decision."""
if not args.llm_merge_decider or not args.post_correct_llm or not llm_state.get("available", True):
return True, candidate_text
prompt = (
"Decide if two consecutive captions should be merged.\n"
"Return JSON only:\n"
"{\"merge\": true|false, \"text\": \"...\"}\n"
"Rules:\n"
"1. Merge only if current line clearly continues previous line.\n"
"2. Never drop unique information from either line.\n"
"3. If uncertain, return {\"merge\": false, \"text\": \"\"}.\n"
"4. If merge=true, output one corrected merged line.\n"
"5. Do not invent content.\n"
f"Previous: {previous_text}\n"
f"Current: {current_text}\n"
f"HeuristicMerged: {candidate_text}\n"
"JSON:"
)
try:
response = call_ollama_generate(args, prompt, args.llm_merge_timeout, temperature=0.0)
response.raise_for_status()
raw = response.json().get("response", "").strip()
if not raw:
return True, candidate_text
try:
parsed = json.loads(raw)
except json.JSONDecodeError:
match = re.search(r"\{.*\}", raw, re.DOTALL)
if not match:
return True, candidate_text
parsed = json.loads(match.group(0))
merge = bool(parsed.get("merge", False))
llm_text = normalize_english_caption(parsed.get("text", "").strip())
if not merge:
return False, ""
if not llm_text:
llm_text = candidate_text
# Safety: reject large rewrites and keep heuristic merge.
if len(llm_text) > max(2 * len(candidate_text), len(candidate_text) + 80):
return True, candidate_text
merged_source = f"{previous_text} {current_text}".strip()
if word_overlap_ratio(merged_source, llm_text) < args.post_correct_min_overlap:
return True, candidate_text
return True, llm_text
except Exception as exc:
if not llm_state.get("merge_warned", False):
print(f"\n[SYSTEM]: LLM merge-decider unavailable ({exc}). Falling back to heuristic merge.")
llm_state["merge_warned"] = True
return True, candidate_text
def normalize_arabic_punctuation(text):
"""Convert common Latin punctuation to Arabic-friendly equivalents."""
updated = text.replace(",", "،").replace(";", "؛")
# Convert question mark only when there are Arabic letters present.
if re.search(r"[\u0600-\u06FF]", updated):
updated = updated.replace("?", "؟")
return updated
def format_caption_lines(lang_key, text):
"""Format caption output lines for terminal display."""
key = lang_key.lower()
if key == "ar":
cleaned = normalize_arabic_punctuation(text)
# Put label on separate line, right-align content, and force RTL embedding.
term_width = shutil.get_terminal_size(fallback=(100, 20)).columns
visible_len = len(cleaned)
left_pad = max(0, term_width - visible_len)
return ["[AR]:", f"{' ' * left_pad}\u202B{cleaned}\u202C"]
return [f"[{lang_key.upper()}]: {text}"]
def print_caption(lang_key, text, leading_newline=False):
lines = format_caption_lines(lang_key, text)
if leading_newline and lines:
lines[0] = "\n" + lines[0]
for line in lines:
print(line)
log_session(line.strip())
def is_hallucination(text):
"""Detect common Whisper hallucinations or high repetition."""
if not text: return False
# Common hallucinations
hallucinations = [
r"thanks? for watching",
r"please subscribe",
r"youtube",
r"click the link",
r"like and subscribe",
r"tuned in",
r"next time",
]
for pattern in hallucinations:
if re.search(pattern, text, re.IGNORECASE):
return True
# Check for excessive word repetition (e.g. "Hallelujah" repeated 10 times)
words = text.lower().split()
if len(words) >= 8:
counts = Counter(words)
most_common_word, count = counts.most_common(1)[0]
if count / len(words) > 0.75:
return True
return False
def callback(indata, frames, time, status):
if status:
print(status, file=sys.stderr)
audio_queue.put(indata.copy())
def ingest_worker():
"""Background thread to handle server ingestion with retries."""
while True:
payload = ingest_queue.get()
if payload is None: break
delay = 1
max_delay = 15
success = False
while not success:
try:
response = requests.post(INGEST_URL, json=payload, timeout=5)
if response.status_code == 200:
success = True
else:
print(f"\n[Ingest Error] Server returned {response.status_code}. Retrying in {delay}s...")
except Exception as e:
print(f"\n[Ingest Error] {e}. Retrying in {delay}s...")
if not success:
time.sleep(delay)
delay = min(delay * 2, max_delay)
ingest_queue.task_done()
def warmup_ollama_model(args):
"""Warm up local Ollama model so first correction request avoids cold-start latency."""
try:
response = call_ollama_generate(
args,
"Return exactly: ok",
timeout=args.post_correct_warmup_timeout,
temperature=0.0,
)
response.raise_for_status()
content = response.json().get("response", "").strip().lower()
if content:
print(f"[SYSTEM]: Ollama warmup OK ({args.post_correct_model}).")
except Exception as exc:
print(f"[SYSTEM]: Ollama warmup failed ({exc}). Continuing; runtime fallback will apply.")
def main():
global WHISPER_MODEL, session_log_handle
parser = argparse.ArgumentParser(description="Live transcription and translation with Whisper.")
parser.add_argument("-es", action="store_true", help="Enable Spanish translation")
parser.add_argument("-en", action="store_true", help="Enable English (detected or bridged)")
parser.add_argument("-ar", action="store_true", help="Enable Arabic translation")
parser.add_argument("-fr", action="store_true", help="Enable French translation")
parser.add_argument("-i", "--ingest", action="store_true", help="Enable data transmission to server")
parser.add_argument("-l", "--list-devices", action="store_true", help="Show available audio devices and exit")
parser.add_argument("-d", "--device", type=int, help="Input device index")
parser.add_argument("--loopback", action="store_true", help="Automatically select a loopback device (e.g., BlackHole, Stereo Mix)")
parser.add_argument("-q", "--quantize", action="store_true", help="Use 4-bit quantized Whisper model for speed")
parser.add_argument("-s", "--stream", action="store_true", help="Enable real-time streaming transcription (Draft mode)")
parser.add_argument("-c", "--context", action="store_true", help="Enable prompt caching/rolling context for better continuity")
parser.add_argument("--lang", type=str, help="Hardcode source language (e.g. 'en', 'es') to bypass detection")
parser.add_argument("--silence", type=int, default=1000, help="Minimum silence duration in ms to end a chunk (default: 1000)")
parser.add_argument("--max-buffer", type=int, default=20, help="Maximum buffer duration in seconds before forcing a flush (default: 20)")
parser.add_argument("--channels", type=int, default=1, help="Number of input channels (default: 1)")
parser.add_argument("--pick-channel", type=int, choices=[0, 1], help="Pick a specific channel (0 or 1) from stereo input")
parser.add_argument("--filter-lang", action="store_true", help="Discard segments where detected language does not match --lang")
parser.add_argument("--draft-beam-size", type=int, default=1, help="Beam size for draft mode transcriptions (default: 1)")
parser.add_argument("--final-beam-size", type=int, default=1, help="Beam size for final segment transcriptions (default: 1)")
parser.add_argument("--temperature-fallback", type=parse_temperature_fallback, default=(0.0, 0.2, 0.4, 0.6, 0.8, 1.0), help="Comma-separated temperatures for fallback decoding (default: 0.0,0.2,0.4,0.6,0.8,1.0)")
parser.add_argument("--logprob-threshold", type=float, default=-0.8, help="Reject low-confidence tokens below this avg logprob (default: -0.8)")
parser.add_argument("--compression-threshold", type=float, default=2.2, help="Reject repetitive outputs above this compression ratio (default: 2.2)")
parser.add_argument("--mt-max-chars", type=int, default=250, help="Max chars per translation chunk before sentence-aware splitting (default: 250)")
parser.add_argument("--mt-max-new-tokens", type=int, default=150, help="Max new tokens per translation chunk (default: 150)")
parser.add_argument("--mt-num-beams", type=int, default=4, help="Beam size for Marian translation generation (default: 4)")
parser.add_argument("--mt-no-repeat-ngram-size", type=int, default=3, help="No-repeat n-gram size for Marian generation (default: 3)")
parser.add_argument("--mt-length-penalty", type=float, default=1.0, help="Length penalty for Marian generation (default: 1.0)")
parser.add_argument("--mt-repetition-penalty", type=float, default=1.05, help="Repetition penalty for Marian generation (default: 1.05)")
parser.add_argument("--mt-no-early-stopping", action="store_true", help="Disable early stopping in Marian beam search")
parser.add_argument("--smart-correct", action="store_true", help="Enable English caption post-correction (retry + glossary + rolling merge)")
parser.add_argument("--retry-logprob-threshold", type=float, default=-1.05, help="Retry transcription when mean avg_logprob is below this threshold (default: -1.05)")
parser.add_argument("--retry-compression-threshold", type=float, default=2.4, help="Retry transcription when max compression ratio exceeds this threshold (default: 2.4)")
parser.add_argument("--caption-correction-window", type=float, default=3.0, help="Seconds where previous English line can still be merged/corrected (default: 3.0)")
parser.add_argument("--glossary-file", type=str, default="glossary.txt", help="Path to glossary file with SOURCE=TARGET pairs (default: glossary.txt if present)")
parser.add_argument("--glossary-pair", action="append", type=parse_glossary_pair, default=[], help="Term replacement pair SOURCE=TARGET (repeatable)")
parser.add_argument("--post-correct", action="store_true", help="Enable v2 finalized English caption post-correction")
parser.add_argument("--post-correct-llm", action="store_true", help="Use local LLM for post-correction (fallback to rules on failure)")
parser.add_argument("--post-correct-model", type=str, default="llama3.1:8b-instruct", help="Local Ollama model for post-correction")
parser.add_argument("--post-correct-ollama-url", type=str, default="http://127.0.0.1:11434/api/generate", help="Ollama generate endpoint for local post-correction")
parser.add_argument("--post-correct-llm-timeout", type=float, default=8.0, help="Timeout seconds for local LLM post-correction (default: 8.0)")
parser.add_argument("--post-correct-keep-alive", type=str, default="30m", help="Ollama keep_alive duration to keep the model loaded (default: 30m)")
parser.add_argument("--post-correct-warmup-timeout", type=float, default=12.0, help="Timeout seconds for startup Ollama warmup (default: 12.0)")
parser.add_argument("--skip-post-correct-warmup", action="store_true", help="Skip startup warmup call for the post-correct LLM")
parser.add_argument("--post-correct-min-overlap", type=float, default=0.45, help="Minimum token overlap ratio required to accept LLM rewrite (default: 0.45)")
parser.add_argument("--post-correct-debug", action="store_true", help="Print pre/post LLM correction text for debugging")
parser.add_argument("--llm-merge-decider", action="store_true", help="Use local LLM to validate/refine smart-correct line merges")
parser.add_argument("--llm-merge-timeout", type=float, default=1.5, help="Timeout seconds for LLM merge decision (default: 1.5)")
parser.add_argument("--speaker-change-detect", action="store_true", help="Enable lightweight speaker-change detection to force caption line cuts")
parser.add_argument("--speaker-sim-threshold", type=float, default=0.72, help="Cosine similarity threshold below which a speaker change is assumed (default: 0.72)")
parser.add_argument("--speaker-min-buffer", type=float, default=1.6, help="Minimum buffered speech seconds before speaker-change checks (default: 1.6)")
parser.add_argument("--speaker-check-interval", type=float, default=0.6, help="Seconds between speaker-change checks while buffering (default: 0.6)")
parser.add_argument("--session-log-file", type=str, default="transcribe_session.log", help="Path to session debug log file (default: transcribe_session.log)")
parser.add_argument("--disable-session-log", action="store_true", help="Disable writing session debug logs to file")
args = parser.parse_args()
if not args.disable_session_log:
session_log_handle = open(args.session_log_file, "a", encoding="utf-8")
log_session("=" * 72)
log_session("SESSION_START")
log_session(f"ARGS {json.dumps(vars(args), sort_keys=True)}")
print(f"Session logging to: {args.session_log_file}")
if args.post_correct_llm and not args.post_correct:
args.post_correct = True
print("[SYSTEM]: Enabling --post-correct because --post-correct-llm was requested.")
log_session("[SYSTEM]: Enabling --post-correct because --post-correct-llm was requested.")
file_glossary_pairs = load_glossary_file(args.glossary_file)
merged_glossary_pairs = file_glossary_pairs + args.glossary_pair
if file_glossary_pairs:
print(f"Loaded {len(file_glossary_pairs)} glossary terms from {args.glossary_file}")
log_session(f"Loaded {len(file_glossary_pairs)} glossary terms from {args.glossary_file}")
if args.post_correct:
print("Post-correct mode: rules enabled.")
log_session("Post-correct mode: rules enabled.")
if args.post_correct_llm:
print(
f"Post-correct LLM: {args.post_correct_model} via {args.post_correct_ollama_url} "
f"(keep_alive={args.post_correct_keep_alive})"
)
log_session(
f"Post-correct LLM: {args.post_correct_model} via {args.post_correct_ollama_url} "
f"(keep_alive={args.post_correct_keep_alive})"
)
if args.post_correct_debug:
print("Post-correct debug logging enabled.")
log_session("Post-correct debug logging enabled.")
if args.llm_merge_decider and not args.post_correct_llm:
print("[SYSTEM]: --llm-merge-decider requires --post-correct-llm. Falling back to heuristic merge.")
args.llm_merge_decider = False
log_session("[SYSTEM]: --llm-merge-decider requires --post-correct-llm. Falling back to heuristic merge.")
if args.speaker_change_detect:
print(
"Speaker-change detection enabled "
f"(threshold={args.speaker_sim_threshold}, min_buffer={args.speaker_min_buffer}s)."
)
log_session(
"Speaker-change detection enabled "
f"(threshold={args.speaker_sim_threshold}, min_buffer={args.speaker_min_buffer}s)."
)
if args.quantize:
WHISPER_MODEL = "mlx-community/whisper-small-mlx-4bit"
log_session("Using quantized Whisper model variant.")
if args.list_devices:
print("\nAvailable Audio Devices:")
print(sd.query_devices())
return
buffer_limit = SAMPLERATE * args.max_buffer
device = "mps" if torch.backends.mps.is_available() else "cpu"
print(f"Using device: {device}")
log_session(f"Using device: {device}")
# Only start ingest thread if enabled
if args.ingest:
threading.Thread(target=ingest_worker, daemon=True).start()
# 1. Load models
print(f"Loading Multilingual Whisper model '{WHISPER_MODEL}'...")
if args.post_correct_llm and not args.skip_post_correct_warmup:
warmup_ollama_model(args)
# Filter translation models to only those enabled by args
active_target_langs = {k: v for k, v in TARGET_LANGS.items() if getattr(args, k, False)}
translation_engines = {}
for lang_key, model_id in active_target_langs.items():
print(f"Loading {lang_key} translation model ({model_id})...")
tokenizer = MarianTokenizer.from_pretrained(model_id)
model = MarianMTModel.from_pretrained(model_id).to(device)
translation_engines[lang_key] = (model, tokenizer)
print("Loading Silero VAD model...")
vad_model = load_silero_vad()
print("Models loaded.")
# 2. Select Audio Device
device_index = None
if args.device is not None:
device_index = args.device
elif args.loopback:
print("Searching for loopback device...")
devices = sd.query_devices()
for i, dev in enumerate(devices):
name = dev['name'].lower()
if any(keyword in name for keyword in ["blackhole", "loopback", "soundflower", "stereo mix"]):
if dev['max_input_channels'] > 0:
device_index = i
print(f"Using loopback device: {dev['name']} (Index {i})")
print("\n[TIP] For macOS with BlackHole:")
print("1. Open 'Audio MIDI Setup' and create a 'Multi-Output Device'.")
print("2. Select your speakers AND 'BlackHole 2ch'.")
print("3. Set your system output to this 'Multi-Output Device'.")
print("This way you can hear the audio while it is being captioned.\n")
break
if device_index is None:
print("No loopback device found. Falling back to default.")
if device_index is None and not args.loopback:
print("\nAvailable Audio Devices:")
print(sd.query_devices())
try:
device_input = input("\nSelect input device index (or press Enter for default): ")
device_index = int(device_input) if device_input.strip() else None
except ValueError:
print("Invalid input, using default device.")
device_index = None
if device_index is not None:
print(f"Using device index: {device_index}")
else:
print("Using system default input device.")
print(f"\nStarting live transcription{' & server ingest' if args.ingest else ''}... (Press Ctrl+C to stop)")
log_session(f"Starting live transcription (ingest={args.ingest}).")
audio_buffer = []
speech_started = False
last_stream_time = time.time()
last_change_time = time.time()
rolling_context = ""
rolling_context_en = ""
last_draft_text = ""
draft_line_active = False
recent_en_lines = deque(maxlen=2)
llm_post_correct_state = {"available": True, "warned": False, "merge_warned": False}
speaker_reference_signature = None
last_speaker_check_time = 0.0
try:
with sd.InputStream(samplerate=SAMPLERATE, channels=args.channels, callback=callback, blocksize=BLOCK_SIZE, device=device_index):
while True:
while not audio_queue.empty():
data = audio_queue.get()
if args.channels > 1:
if args.pick_channel is not None:
# Select specific channel
data = data[:, args.pick_channel]
else:
# Mix to mono
data = np.mean(data, axis=1)
audio_buffer.append(data.flatten())
if len(audio_buffer) > 0:
current_audio = np.concatenate(audio_buffer)
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
)
# --- STUCK WATCHDOG ---
# If buffer is getting long (>12s) and we haven't had a change in draft for 7s,
# OR if buffer is extremely long (>25s) regardless of draft activity.
time_since_last_change = time.time() - last_change_time
if (buffer_duration > 12.0 and time_since_last_change > 7.0) or (buffer_duration > 25.0):
print(f"\n[SYSTEM]: Transcription watchdog triggered (Buffer: {buffer_duration:.1f}s, No change: {time_since_last_change:.1f}s). Resetting...")
audio_buffer = []
speech_started = False
rolling_context = ""
rolling_context_en = ""
recent_en_lines.clear()
speaker_reference_signature = None
last_draft_text = ""
last_change_time = time.time()
if args.stream:
sys.stdout.write("\r\033[K")
sys.stdout.flush()
continue
if len(speech_timestamps) > 0:
speech_started = True
last_end = speech_timestamps[-1]['end']
buffer_len_samples = len(current_audio)
force_flush_for_speaker_change = False
if args.speaker_change_detect and buffer_duration >= args.speaker_min_buffer:
now = time.time()
if (now - last_speaker_check_time) >= args.speaker_check_interval:
last_speaker_check_time = now
tail_secs = min(2.0, buffer_duration)
tail_samples = int(tail_secs * SAMPLERATE)
recent_signature = compute_speaker_signature(current_audio[-tail_samples:], SAMPLERATE)
if speaker_reference_signature is None and recent_signature is not None:
speaker_reference_signature = recent_signature
elif recent_signature is not None:
similarity = cosine_similarity(speaker_reference_signature, recent_signature)
if similarity is not None and similarity < args.speaker_sim_threshold:
force_flush_for_speaker_change = True
print(
f"\n[SYSTEM]: Speaker change detected (similarity={similarity:.2f}). "
"Cutting current line."
)
if (
(buffer_len_samples - last_end) > (SAMPLERATE * args.silence / 1000)
or buffer_len_samples > buffer_limit
or force_flush_for_speaker_change
):
# Clear draft line if it was used
if args.stream:
sys.stdout.write("\r\033[K")
sys.stdout.flush()
draft_line_active = False
# Prepare transcription kwargs
transcribe_kwargs = {"path_or_hf_repo": WHISPER_MODEL}
if args.lang:
transcribe_kwargs["language"] = args.lang
if args.context and rolling_context:
transcribe_kwargs["initial_prompt"] = rolling_context
if args.final_beam_size > 1:
transcribe_kwargs["beam_size"] = args.final_beam_size
transcribe_kwargs["temperature"] = args.temperature_fallback
transcribe_kwargs["logprob_threshold"] = args.logprob_threshold
transcribe_kwargs["compression_ratio_threshold"] = args.compression_threshold
# 1. Transcribe & Detect Language
transcription_result = transcribe_with_controls(current_audio, transcribe_kwargs)
original_text = transcription_result['text'].strip()
detected_lang = transcription_result.get('language', args.lang if args.lang else 'en')
avg_logprob, max_compression = extract_whisper_quality(transcription_result)
if args.smart_correct and should_retry_transcription(original_text, avg_logprob, max_compression, args):
retry_kwargs = dict(transcribe_kwargs)
retry_kwargs["temperature"] = (0.0,)
retry_kwargs.pop("initial_prompt", None)
retry_result = transcribe_with_controls(current_audio, retry_kwargs)
retry_text = retry_result.get("text", "").strip()
retry_avg_logprob, retry_max_compression = extract_whisper_quality(retry_result)
original_score = quality_score(avg_logprob, max_compression)
retry_score = quality_score(retry_avg_logprob, retry_max_compression)
if retry_text and retry_score > (original_score + 0.05):
print("\n[SYSTEM]: Low-confidence segment re-decoded with better quality.")
transcription_result = retry_result
original_text = retry_text
avg_logprob, max_compression = retry_avg_logprob, retry_max_compression
# Filter language if requested
if args.filter_lang and args.lang and detected_lang != args.lang:
print(f"\n[SYSTEM]: Discarding segment (Detected: {detected_lang}, Expected: {args.lang})")
original_text = ""
if is_hallucination(original_text):
print(f"\n[SYSTEM]: Hallucination detected, ignoring and resetting context.")
original_text = ""
rolling_context = ""
rolling_context_en = ""
recent_en_lines.clear()
speaker_reference_signature = None
if original_text:
print_caption(detected_lang, original_text, leading_newline=True)
last_change_time = time.time() # Successfully transcribed full segment
# Prepare payload
payload = {"original": original_text}
# ... (payload construction)
# Include detected language if requested or if it's the bridge
if (detected_lang in active_target_langs) or (detected_lang == "en" and args.en):
payload[detected_lang] = original_text
# 2. Bridge to English if not already English
if detected_lang != "en":
bridge_kwargs = {"path_or_hf_repo": WHISPER_MODEL, "task": "translate"}
if args.lang:
bridge_kwargs["language"] = args.lang
if args.context and rolling_context_en:
# Keep a dedicated English context for translate mode.
bridge_kwargs["initial_prompt"] = rolling_context_en
if args.final_beam_size > 1:
bridge_kwargs["beam_size"] = args.final_beam_size
bridge_kwargs["temperature"] = args.temperature_fallback
bridge_kwargs["logprob_threshold"] = args.logprob_threshold
bridge_kwargs["compression_ratio_threshold"] = args.compression_threshold
bridge_result = transcribe_with_controls(current_audio, bridge_kwargs)
english_text = bridge_result['text'].strip()
if args.context and english_text:
rolling_context_en = (rolling_context_en + " " + english_text)[-200:].strip()
else:
english_text = original_text
if args.context and english_text:
rolling_context_en = (rolling_context_en + " " + english_text)[-200:].strip()
english_for_translation = english_text
if args.smart_correct and english_for_translation:
english_for_translation = normalize_english_caption(english_for_translation)
english_for_translation = apply_glossary(english_for_translation, merged_glossary_pairs)
elif merged_glossary_pairs and english_for_translation:
# Allow glossary use even when smart-correct is disabled.
english_for_translation = apply_glossary(english_for_translation, merged_glossary_pairs)
if args.post_correct and english_for_translation:
context_lines = [item["text"] for item in recent_en_lines]
english_for_translation = apply_rule_based_post_correction(english_for_translation)
english_for_translation = post_correct_with_local_llm(
english_for_translation,
args,
llm_post_correct_state,
context_lines=context_lines,
)
english_for_caption = english_for_translation
if args.smart_correct and english_for_caption:
now_ts = time.time()
merged = maybe_merge_recent_caption(
recent_en_lines,
english_for_caption,
now_ts,
args.caption_correction_window,
)
if merged:
old_line, new_line = merged
merge_ok, merge_text = llm_decide_caption_merge(
old_line,
english_for_caption,
new_line,
args,
llm_post_correct_state,
)
if merge_ok:
english_for_caption = merge_text
print(f"[EN-REV]: {old_line} -> {english_for_caption}")
log_session(f"[EN-REV]: {old_line} -> {english_for_caption}")
else:
if recent_en_lines:
recent_en_lines[-1]["text"] = old_line
recent_en_lines[-1]["ts"] = now_ts
recent_en_lines.append({"text": english_for_caption, "ts": now_ts})
else:
recent_en_lines.append({"text": english_for_caption, "ts": now_ts})
if args.en and english_for_caption:
payload["en"] = english_for_caption
print(f"[EN]: {english_for_caption}")
log_session(f"[EN]: {english_for_caption}")
# Update rolling context for next segment
if args.context:
# keep the last ~200 characters of the source language text
rolling_context = (rolling_context + " " + original_text)[-200:].strip()
# 3. Translate from English to other languages
if english_for_translation and translation_engines:
english_chunks = split_text_for_translation(english_for_translation, max_chars=args.mt_max_chars)
for lang_key, (model, tokenizer) in translation_engines.items():
# Skip if we already filled this (e.g. detected lang was 'es')
if lang_key in payload:
if lang_key != detected_lang: # Already printed original
print_caption(lang_key, payload[lang_key])
continue
translated_parts = []
for chunk in english_chunks:
inputs = tokenizer(chunk, return_tensors="pt", padding=True).to(device)
with torch.no_grad():
translated_tokens = model.generate(
**inputs,
max_new_tokens=args.mt_max_new_tokens,
num_beams=args.mt_num_beams,
no_repeat_ngram_size=args.mt_no_repeat_ngram_size,
length_penalty=args.mt_length_penalty,
repetition_penalty=args.mt_repetition_penalty,
early_stopping=not args.mt_no_early_stopping,
)
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_caption(lang_key, translated_text)
# Queue for background ingestion if enabled
if args.ingest:
ingest_queue.put(payload)
# print(f"Sent to ingest: {list(payload.keys())}")
audio_buffer = []
speech_started = False
last_stream_time = time.time()
stuck_draft_count = 0
last_draft_text = ""
speaker_reference_signature = None
draft_line_active = False
elif args.stream and (time.time() - last_stream_time) > 1.0:
# Draft transcription
draft_kwargs = {"path_or_hf_repo": WHISPER_MODEL}
if args.lang: draft_kwargs["language"] = args.lang
if args.context and rolling_context: draft_kwargs["initial_prompt"] = rolling_context
if args.draft_beam_size > 1:
draft_kwargs["beam_size"] = args.draft_beam_size
draft_kwargs["temperature"] = 0.0
draft_result = transcribe_with_controls(current_audio, draft_kwargs)
draft_text = draft_result['text'].strip()
draft_lang = draft_result.get('language', args.lang if args.lang else 'en')
if args.filter_lang and args.lang and draft_lang != args.lang:
draft_text = ""
if draft_text:
# Update timestamp ONLY if the text actually changed
if draft_text != last_draft_text:
last_change_time = time.time()
last_draft_text = draft_text
if not draft_line_active:
# Reserve one blank line before the rolling draft line.
sys.stdout.write("\n")
draft_line_active = True
sys.stdout.write(f"\r\033[K[DRAFT]: {draft_text}")
sys.stdout.flush()
# Send draft to ingest if enabled
if args.ingest:
ingest_queue.put({"draft": draft_text})
log_session("[DRAFT]: queued draft ingest payload")
else:
# If Whisper returns empty, check if we've been silent for too long
# even though VAD says there is speech.
if buffer_duration > 10.0 and (time.time() - last_change_time) > 7.0:
print(f"\n[SYSTEM]: Draft is empty while audio continues. Forcing reset...")
audio_buffer = []
speech_started = False
rolling_context = ""
rolling_context_en = ""
recent_en_lines.clear()
speaker_reference_signature = None
last_change_time = time.time()
last_draft_text = ""
sys.stdout.write("\r\033[K")
sys.stdout.flush()
draft_line_active = False
continue
last_stream_time = time.time()
elif not speech_started and len(current_audio) > SAMPLERATE * 2:
audio_buffer = []
speaker_reference_signature = None
except KeyboardInterrupt:
print("\nStopped by user.")
log_session("Stopped by user.")
except Exception as e:
print(f"\nError: {e}")
log_session(f"Error: {e}")
finally:
if session_log_handle is not None:
log_session("SESSION_END")
session_log_handle.close()
session_log_handle = None
if __name__ == "__main__":
import multiprocessing
multiprocessing.freeze_support()
main()