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 from pyannote.audio import Pipeline from dotenv import load_dotenv from langdetect import detect, DetectorFactory from langdetect.lang_detect_exception import LangDetectException # Enforce deterministic langdetect results DetectorFactory.seed = 0 # Load environment variables from .env file load_dotenv() # 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 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]: ", 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 structure_paragraph_with_llm(text, args, llm_state): """Use a local LLM to structure a series of captions into a paragraph.""" prompt = ( "You are a real-time English caption editor. Your task is to take a series of transcribed audio segments and format them into a single, coherent paragraph.\n" "Rules:\n" "1. Combine the segments into a flowing paragraph.\n" "2. Correct grammar, punctuation, and casing.\n" "3. Remove disfluencies and false starts (e.g., 'um', 'uh').\n" "4. Do not change the meaning of the text.\n" "5. If there are speaker labels (e.g., [SPEAKER_01]:), preserve them.\n" "6. Output ONLY the formatted paragraph.\n\n" "Here are the segments:\n" f"{text}\n\n" "Formatted paragraph:" ) try: response = call_ollama_generate(args, prompt, args.post_correct_llm_timeout, temperature=0.2) response.raise_for_status() raw = response.json().get("response", "").strip() return raw except Exception as exc: if not llm_state.get("paragraph_warned", False): print(f"\n[SYSTEM]: Local paragraph structuring LLM unavailable ({exc}).") llm_state["paragraph_warned"] = True return text # Fallback to the original 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 assign_speakers_to_segments(segments, diarization): """Assign speaker labels to transcription segments based on diarization.""" if not diarization or not segments: return segments for segment in segments: segment_start = segment['start'] segment_end = 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: intersection_duration = intersection_end - intersection_start if speaker not in speaker_intersections: speaker_intersections[speaker] = 0 speaker_intersections[speaker] += intersection_duration if speaker_intersections: dominant_speaker = max(speaker_intersections, key=speaker_intersections.get) segment['speaker'] = dominant_speaker else: # Keep speaker from previous segment if segment is short and follows closely # This can help with short utterances that are missed by the diarizer # For now, we'll just label as unknown segment['speaker'] = 'UNKNOWN' return segments 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("--detect-auto-translation", action="store_true", help="Enable heuristic to detect and leverage auto-translated English segments.") 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="qwen:2b", 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-paragraph", action="store_true", help="Enable LLM-based paragraph structuring.") 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-diarization", action="store_true", help="Enable speaker diarization using pyannote.audio.") 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.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}") diarization_pipeline = None if args.speaker_diarization: print("Loading speaker diarization pipeline...") hf_token = os.environ.get("HF_TOKEN") if not hf_token: print("[WARN] Hugging Face token not found (HF_TOKEN env variable). Diarization may fail if model is private.") log_session("[WARN] Hugging Face token not found (HF_TOKEN env variable).") try: diarization_pipeline = Pipeline.from_pretrained( "pyannote/speaker-diarization-3.1" ) diarization_pipeline.to(torch.device(device)) print("Speaker diarization pipeline loaded.") log_session("Speaker diarization pipeline loaded.") except Exception as e: print(f"[ERROR] Could not load speaker diarization pipeline: {e}") log_session(f"[ERROR] Could not load speaker diarization pipeline: {e}") diarization_pipeline = None # 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, "paragraph_warned": False} english_paragraph_buffer = [] 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() 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) if ( (buffer_len_samples - last_end) > (SAMPLERATE * args.silence / 1000) or buffer_len_samples > buffer_limit ): # 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, "word_timestamps": args.speaker_diarization} 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 # Speaker diarization diarization_result = None if diarization_pipeline: log_session("Performing speaker diarization...") audio_for_diarization = torch.from_numpy(current_audio).float().unsqueeze(0) try: diarization_result = diarization_pipeline({"waveform": audio_for_diarization, "sample_rate": SAMPLERATE}) except Exception as e: print(f"\n[ERROR] Diarization failed: {e}") log_session(f"Diarization failed: {e}") # 1. Transcribe & Detect Language transcription_result = transcribe_with_controls(current_audio, transcribe_kwargs) original_text_for_print = "" if diarization_result: segments_with_speakers = assign_speakers_to_segments( transcription_result.get('segments', []), diarization_result ) transcription_result['segments'] = segments_with_speakers # Reconstruct text with speaker labels for printing text_parts_for_print = [] for i, segment in enumerate(segments_with_speakers): speaker = segment.get('speaker', 'UNKNOWN') text = segment['text'].strip() if i > 0 and speaker == segments_with_speakers[i-1].get('speaker', 'UNKNOWN'): text_parts_for_print.append(f" {text}") else: if i > 0: text_parts_for_print.append("\n") text_parts_for_print.append(f"[{speaker}]: {text}") original_text_for_print = "".join(text_parts_for_print) # Reconstruct text without speaker labels for processing original_text = " ".join([seg['text'].strip() for seg in transcription_result['segments']]) else: original_text = transcription_result['text'].strip() original_text_for_print = original_text 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() if original_text: if args.detect_auto_translation and args.lang and args.lang != 'en' and len(original_text.split()) > 3: try: detected_lang_for_segment = detect(original_text) if detected_lang_for_segment == 'en': print_caption("AUTO-EN", original_text, leading_newline=True) if args.context: rolling_context_en = (rolling_context_en + " " + original_text)[-200:].strip() except LangDetectException: # Could not detect language, probably too short pass print_caption(detected_lang, original_text_for_print, 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.llm_paragraph: english_paragraph_buffer.append(english_for_caption) if english_for_caption.strip().endswith((".", "?", "!")): full_text = "\n".join(english_paragraph_buffer) structured_paragraph = structure_paragraph_with_llm(full_text, args, llm_post_correct_state) print(f"\n[LLM Paragraph]:\n{structured_paragraph}") log_session(f"[LLM-P]: {structured_paragraph}") # Translate the paragraph here if structured_paragraph and translation_engines: english_chunks = split_text_for_translation(structured_paragraph, max_chars=args.mt_max_chars) for lang_key, (model, tokenizer) in translation_engines.items(): 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() print_caption(lang_key, translated_text) english_paragraph_buffer = [] # Reset buffer if args.en and english_for_caption: payload["en"] = english_for_caption # When llm_paragraph is enabled, we don't print the line-by-line English captions. if not args.llm_paragraph: 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 not args.llm_paragraph and 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 = "" 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 = [] 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()