diff --git a/config.json b/config.json new file mode 100644 index 0000000..1fb7482 --- /dev/null +++ b/config.json @@ -0,0 +1,12 @@ +{ + "model": "mlx-community/whisper-base-mlx", + "device": null, + "lang": null, + "silence": 1000, + "max_buffer": 20, + "channels": 1, + "es": false, + "fr": false, + "ar": false, + "ingest": false +} diff --git a/engine_distribute.py b/engine_distribute.py new file mode 100644 index 0000000..69f4a83 --- /dev/null +++ b/engine_distribute.py @@ -0,0 +1,49 @@ +import requests +import time +import argparse +import multiprocessing + +INGEST_URL = "https://emiapi.reynafamily.com/live-captions/ingest" + +def run_distribution(in_queue, args): + print("[Distribute] Starting distribution service...") + + while True: + try: + payload = in_queue.get() + if payload is None: break + + if args.ingest: + 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"[Distribute Error] {response.status_code}. Retrying in {delay}s...") + except Exception as e: + print(f"[Distribute Error] {e}. Retrying in {delay}s...") + + if not success: + time.sleep(delay) + delay = min(delay * 2, max_delay) + + print(f"[Distribute] Sent payload: {list(payload.keys())}") + else: + print(f"[Distribute] Local display (ingest disabled): {payload}") + + except Exception as e: + print(f"[Distribute] Error: {e}") + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("-i", "--ingest", action="store_true") + args = parser.parse_args() + + # Dummy queue for testing + in_q = multiprocessing.Queue() + run_distribution(in_q, args) diff --git a/engine_transcribe.py b/engine_transcribe.py new file mode 100644 index 0000000..6aa1172 --- /dev/null +++ b/engine_transcribe.py @@ -0,0 +1,143 @@ +import sys +from unittest.mock import MagicMock + +# 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 sys +import time +import numpy as np +import sounddevice as sd +import torch +import mlx_whisper +from silero_vad import load_silero_vad, get_speech_timestamps +import queue +import argparse + +# Audio Constants +SAMPLERATE = 16000 +BLOCK_SIZE = 512 +VAD_THRESHOLD = 0.5 + +def transcribe_with_controls(audio, transcribe_kwargs): + """Call mlx_whisper.transcribe with graceful fallback for unsupported args.""" + try: + return mlx_whisper.transcribe(audio, **transcribe_kwargs) + except Exception as exc: + # Fallback logic for beam_size or other specific MLX implementation gaps + if "beam_size" in transcribe_kwargs: + transcribe_kwargs.pop("beam_size") + return mlx_whisper.transcribe(audio, **transcribe_kwargs) + raise exc + +def audio_callback(indata, frames, time, status, audio_queue): + if status: + print(status, file=sys.stderr) + audio_queue.put(indata.copy()) + +def run_transcription(out_queue, args): + device = "mps" if torch.backends.mps.is_available() else "cpu" + print(f"[Transcribe] Loading Whisper model '{args.model}' on {device}...") + + vad_model = load_silero_vad() + audio_queue = queue.Queue() + + # Selection of device + device_index = args.device + + def callback_wrapper(indata, frames, time, status): + audio_callback(indata, frames, time, status, audio_queue) + + audio_buffer = [] + speech_started = False + buffer_limit = SAMPLERATE * args.max_buffer + + print(f"[Transcribe] Starting audio stream on device {device_index}...") + + try: + with sd.InputStream(samplerate=SAMPLERATE, channels=args.channels, callback=callback_wrapper, blocksize=BLOCK_SIZE, device=device_index): + while True: + while not audio_queue.empty(): + data = audio_queue.get() + if args.channels > 1: + data = np.mean(data, axis=1) # Mix to mono + audio_buffer.append(data.flatten()) + + if audio_buffer: + 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 + ) + + if speech_timestamps: + 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: + # 1. Transcribe + transcribe_kwargs = { + "path_or_hf_repo": args.model, + "temperature": (0.0, 0.2, 0.4, 0.6, 0.8, 1.0), + } + if args.lang: + transcribe_kwargs["language"] = args.lang + + result = transcribe_with_controls(current_audio, transcribe_kwargs) + text = result['text'].strip() + detected_lang = result.get('language', 'en') + + if text: + # Send to Translation Process + out_queue.put({ + "original": text, + "detected_lang": detected_lang, + "ts": time.time() + }) + print(f"[Transcribe] {detected_lang.upper()}: {text}") + + audio_buffer = [] + speech_started = False + + elif not speech_started and len(current_audio) > SAMPLERATE * 2: + audio_buffer = [] + + except KeyboardInterrupt: + print("[Transcribe] Stopped.") + except Exception as e: + print(f"[Transcribe] Error: {e}") + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--model", type=str, default="mlx-community/whisper-base-mlx") + parser.add_argument("--device", type=int, default=None) + parser.add_argument("--lang", type=str, default=None) + parser.add_argument("--silence", type=int, default=1000) + parser.add_argument("--max-buffer", type=int, default=20) + parser.add_argument("--channels", type=int, default=1) + args = parser.parse_args() + + # This part would normally be called by the main coordinator + # For testing, we can use a dummy queue + import multiprocessing + q = multiprocessing.Queue() + run_transcription(q, args) diff --git a/engine_translate.py b/engine_translate.py new file mode 100644 index 0000000..d24fc7e --- /dev/null +++ b/engine_translate.py @@ -0,0 +1,130 @@ +import sys +from unittest.mock import MagicMock + +# 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 torch +from transformers import MarianMTModel, MarianTokenizer +import re +import argparse +import multiprocessing +import time + +TARGET_LANGS = { + "es": "Helsinki-NLP/opus-mt-en-es", + "fr": "Helsinki-NLP/opus-mt-en-fr", + "ar": "Helsinki-NLP/opus-mt-en-ar" +} + +def split_text_for_translation(text, max_chars=250): + normalized = " ".join(text.split()).strip() + if not normalized or len(normalized) <= max_chars: + return [normalized] if normalized else [] + + chunks = [] + current = "" + sentences = [s for s in re.split(r"(?<=[.!?])\s+", normalized) if s] + + for sentence in sentences: + if len(sentence) > max_chars: + words = sentence.split() + word_chunk = "" + for word in words: + candidate = f"{word_chunk} {word}".strip() + if len(candidate) <= max_chars: + word_chunk = candidate + 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 run_translation(in_queue, out_queue, args): + device = "mps" if torch.backends.mps.is_available() else "cpu" + print(f"[Translate] Loading translation models on {device}...") + + translation_engines = {} + for lang_key, model_id in TARGET_LANGS.items(): + if getattr(args, lang_key, False): + print(f"[Translate] Loading {lang_key} model...") + tokenizer = MarianTokenizer.from_pretrained(model_id) + model = MarianMTModel.from_pretrained(model_id).to(device) + translation_engines[lang_key] = (model, tokenizer) + + paragraph_buffer = [] + + print("[Translate] Ready.") + + while True: + try: + item = in_queue.get() + if item is None: break + + text = item.get("original", "").strip() + detected_lang = item.get("detected_lang", "en") + + # Simple paragraph logic: + # If the segment ends with sentence-terminal punctuation, flush the paragraph. + paragraph_buffer.append(text) + + if text.endswith((".", "?", "!")): + full_text = " ".join(paragraph_buffer) + payload = {"original": full_text, "en": full_text if detected_lang == "en" else None} + + # If detected language is not English, we'd normally bridge to English first. + # For now, let's assume direct translation for simplicity or bridge if needed. + # (Refining bridge logic can come later) + + for lang_key, (model, tokenizer) in translation_engines.items(): + chunks = split_text_for_translation(full_text) + translated_parts = [] + for chunk in chunks: + inputs = tokenizer(chunk, return_tensors="pt", padding=True).to(device) + with torch.no_grad(): + translated_tokens = model.generate(**inputs, max_new_tokens=150) + translated_parts.append(tokenizer.decode(translated_tokens[0], skip_special_tokens=True).strip()) + + translated_text = " ".join(part for part in translated_parts if part).strip() + payload[lang_key] = translated_text + print(f"[Translate] {lang_key.upper()}: {translated_text}") + + out_queue.put(payload) + paragraph_buffer = [] + + except Exception as e: + print(f"[Translate] Error: {e}") + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("-es", action="store_true") + parser.add_argument("-fr", action="store_true") + parser.add_argument("-ar", action="store_true") + args = parser.parse_args() + + # Dummy queues for testing + in_q = multiprocessing.Queue() + out_q = multiprocessing.Queue() + run_translation(in_q, out_q, args) diff --git a/main_v2.py b/main_v2.py new file mode 100644 index 0000000..c14ac99 --- /dev/null +++ b/main_v2.py @@ -0,0 +1,84 @@ +import sys +from unittest.mock import MagicMock + +# 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 multiprocessing +import argparse +import sys +import json +import os +from engine_transcribe import run_transcription +from engine_translate import run_translation +from engine_distribute import run_distribution + +def main(): + config_path = "config.json" + defaults = {} + if os.path.exists(config_path): + with open(config_path, "r") as f: + defaults = json.load(f) + + parser = argparse.ArgumentParser(description="Multi-process Transcription and Translation.") + + # Transcribe Args + parser.add_argument("--model", type=str, default=defaults.get("model", "mlx-community/whisper-base-mlx")) + parser.add_argument("--device", type=int, default=defaults.get("device")) + parser.add_argument("--lang", type=str, default=defaults.get("lang")) + parser.add_argument("--silence", type=int, default=defaults.get("silence", 1000)) + parser.add_argument("--max-buffer", type=int, default=defaults.get("max_buffer", 20)) + parser.add_argument("--channels", type=int, default=defaults.get("channels", 1)) + + # Translate Args + parser.add_argument("-es", action="store_true", default=defaults.get("es", False), help="Enable Spanish translation") + parser.add_argument("-fr", action="store_true", default=defaults.get("fr", False), help="Enable French translation") + parser.add_argument("-ar", action="store_true", default=defaults.get("ar", False), help="Enable Arabic translation") + + # Distribute Args + parser.add_argument("-i", "--ingest", action="store_true", default=defaults.get("ingest", False), help="Enable data transmission to server") + + args = parser.parse_args() + + # Queues for communication + # Transcribe -> Translate + q_trans_to_tl = multiprocessing.Queue() + # Translate -> Distribute + q_tl_to_dist = multiprocessing.Queue() + + # Processes + p_transcribe = multiprocessing.Process(target=run_transcription, args=(q_trans_to_tl, args)) + p_translate = multiprocessing.Process(target=run_translation, args=(q_trans_to_tl, q_tl_to_dist, args)) + p_distribute = multiprocessing.Process(target=run_distribution, args=(q_tl_to_dist, args)) + + print("[Main] Starting processes...") + p_transcribe.start() + p_translate.start() + p_distribute.start() + + try: + p_transcribe.join() + p_translate.join() + p_distribute.join() + except KeyboardInterrupt: + print("\n[Main] Stopping processes...") + p_transcribe.terminate() + p_translate.terminate() + p_distribute.terminate() + sys.exit(0) + +if __name__ == "__main__": + multiprocessing.freeze_support() + main() diff --git a/requirements.txt b/requirements.txt index fabf9c5..d94fe0a 100644 --- a/requirements.txt +++ b/requirements.txt @@ -9,3 +9,7 @@ pyinstaller sacremoses joblib sentencepiece +pyannote.audio +soundfile +langdetect +python-dotenv diff --git a/transcribe.py b/transcribe.py index d689cb0..5ef2fa1 100644 --- a/transcribe.py +++ b/transcribe.py @@ -33,6 +33,18 @@ 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" @@ -168,51 +180,6 @@ def quality_score(mean_logprob, max_compression): 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 @@ -373,6 +340,34 @@ def post_correct_with_local_llm(text, args, llm_state, context_lines=None): 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 @@ -507,6 +502,38 @@ def is_hallucination(text): 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) @@ -569,6 +596,7 @@ def main(): 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)") @@ -594,7 +622,7 @@ def main(): 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-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)") @@ -602,12 +630,10 @@ def main(): 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-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("--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") @@ -646,15 +672,6 @@ def main(): 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" @@ -670,6 +687,26 @@ def main(): 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() @@ -743,9 +780,8 @@ def main(): 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 + 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): @@ -785,7 +821,6 @@ def main(): rolling_context = "" rolling_context_en = "" recent_en_lines.clear() - speaker_reference_signature = None last_draft_text = "" last_change_time = time.time() if args.stream: @@ -797,30 +832,10 @@ def main(): 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 @@ -830,7 +845,7 @@ def main(): draft_line_active = False # Prepare transcription kwargs - transcribe_kwargs = {"path_or_hf_repo": WHISPER_MODEL} + 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: @@ -841,9 +856,46 @@ def main(): 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 = transcription_result['text'].strip() + + 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) @@ -873,10 +925,20 @@ def main(): 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) + 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 @@ -957,10 +1019,44 @@ def main(): 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 - print(f"[EN]: {english_for_caption}") - log_session(f"[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: @@ -968,7 +1064,7 @@ def main(): rolling_context = (rolling_context + " " + original_text)[-200:].strip() # 3. Translate from English to other languages - if english_for_translation and translation_engines: + 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(): @@ -1007,7 +1103,6 @@ def main(): 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: @@ -1064,7 +1159,6 @@ def main(): elif not speech_started and len(current_audio) > SAMPLERATE * 2: audio_buffer = [] - speaker_reference_signature = None except KeyboardInterrupt: print("\nStopped by user.")