From 5f2738c49c764935ee26130d87435379cdc1c1b6 Mon Sep 17 00:00:00 2001 From: Adolfo Reyna Date: Sun, 15 Mar 2026 19:25:39 -0400 Subject: [PATCH] feat: add speaker diarization and LLM paragraph structuring --- .gitignore | 2 + requirements.txt | 4 + transcribe.py | 278 +++++++++++++++++++++++++++++++---------------- 3 files changed, 192 insertions(+), 92 deletions(-) diff --git a/.gitignore b/.gitignore index 8060f19..548ed60 100644 --- a/.gitignore +++ b/.gitignore @@ -5,3 +5,5 @@ build/ dist/ *.spec venv/ +.env +*.log 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.")