import sys import time import re import os 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-small-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() 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 post_correct_with_local_llm(text, args, llm_state, context_hint=""): """Optionally rewrite finalized English caption via a local Ollama model.""" if not args.post_correct_llm or not llm_state.get("available", True): return text prompt = ( "You correct live English captions.\n" "Rules:\n" "- Preserve original meaning exactly.\n" "- Fix grammar, punctuation, and natural wording.\n" "- Remove filler words only when they do not change meaning.\n" "- Keep output concise and in one line.\n" "- Return only corrected text, no quotes or explanation.\n" ) if context_hint: prompt += f"Previous caption context: {context_hint}\n" prompt += f"Caption: {text}\nCorrected:" try: response = requests.post( args.post_correct_ollama_url, json={ "model": args.post_correct_model, "prompt": prompt, "stream": False, "options": {"temperature": 0.1}, }, timeout=args.post_correct_llm_timeout, ) response.raise_for_status() raw = response.json().get("response", "").strip() cleaned = normalize_english_caption(raw) if cleaned: return cleaned 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 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 = ( "You decide whether two live English caption lines should be merged.\n" "Return strict JSON only in this shape:\n" "{\"merge\": true|false, \"text\": \"final merged text or empty\"}\n" "Rules:\n" "- Preserve meaning exactly.\n" "- If uncertain, return merge=false.\n" "- If merge=true, provide a single corrected line in text.\n" f"Previous: {previous_text}\n" f"Current: {current_text}\n" f"HeuristicMerged: {candidate_text}\n" "JSON:" ) try: response = requests.post( args.post_correct_ollama_url, json={ "model": args.post_correct_model, "prompt": prompt, "stream": False, "options": {"temperature": 0.0}, }, timeout=args.llm_merge_timeout, ) 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 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 and force RTL embedding for the content line. return ["[AR]:", f"\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) 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 main(): global WHISPER_MODEL 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=3.0, help="Timeout seconds for local LLM post-correction (default: 3.0)") 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=0.7, help="Timeout seconds for LLM merge decision (default: 0.7)") args = parser.parse_args() 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.") 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}") if args.post_correct: print("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}") 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 if args.quantize: WHISPER_MODEL = "mlx-community/whisper-small-mlx-4bit" 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}") # 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}'...") # 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)") audio_buffer = [] speech_started = False last_stream_time = time.time() last_change_time = time.time() rolling_context = "" rolling_context_en = "" last_draft_text = "" recent_en_lines = deque(maxlen=2) llm_post_correct_state = {"available": True, "warned": False, "merge_warned": False} 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() # 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() 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_hint = recent_en_lines[-1]["text"] if recent_en_lines else "" 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_hint=context_hint, ) 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}") 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}") # 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 = "" 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 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}) 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() last_change_time = time.time() last_draft_text = "" sys.stdout.write("\r\033[K") sys.stdout.flush() continue last_stream_time = time.time() elif not speech_started and len(current_audio) > SAMPLERATE * 2: audio_buffer = [] except KeyboardInterrupt: print("\nStopped by user.") except Exception as e: print(f"\nError: {e}") if __name__ == "__main__": import multiprocessing multiprocessing.freeze_support() main()