feat: add intelligent english caption correction v1

history.md

@@ -136,3 +136,11 @@ The project now features a high-performance, Apple Silicon-optimized pipeline th
     - **Sentence-Aware MT Chunking:** Replaced hard truncation with sentence-aware splitting for long English bridge text before translation.
     - **Marian Generation Tuning:** Added configurable generation controls (`num_beams`, `no_repeat_ngram_size`, `length_penalty`, `repetition_penalty`, `early_stopping`) to reduce repetitive or unstable outputs.
 - **Outcome:** Better continuity for non-English to English bridging, fewer clipped translations on long segments, and cleaner target-language output with safer defaults for current `mlx_whisper` builds.
+
+## Phase 17: Intelligent English Caption Correction (V1)
+- **Goal:** Improve English caption readability and resilience without introducing heavy latency.
+- **Approach:**
+    - **Confidence-Triggered Re-Decode:** Added optional retry logic that re-transcribes low-confidence segments when Whisper quality signals indicate likely errors (`avg_logprob` / `compression_ratio` thresholds).
+    - **Glossary Replacements:** Added configurable `SOURCE=TARGET` term normalization (`--glossary-pair`) to consistently correct names/terms in English captions.
+    - **Rolling Caption Revision Buffer:** Added a short correction window (`--caption-correction-window`, default 3s) with a 2-line rolling buffer that can merge incomplete English lines into a cleaner sentence as new context arrives.
+- **Outcome:** Captions now support lightweight post-correction behavior that improves continuity and term consistency while remaining compatible with real-time streaming.

transcribe.py

@@ -6,7 +6,7 @@ import threading
 import json
 import argparse
 from unittest.mock import MagicMock
-from collections import Counter
+from collections import Counter, deque
 
 # Comprehensive workaround for missing _lzma in some Python builds
 try:
@@ -128,6 +128,89 @@ def transcribe_with_controls(audio, transcribe_kwargs):
             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 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 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 is_hallucination(text):
     """Detect common Whisper hallucinations or high repetition."""
     if not text: return False
@@ -219,6 +302,11 @@ def main():
     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-pair", action="append", type=parse_glossary_pair, default=[], help="Term replacement pair SOURCE=TARGET (repeatable)")
     
     args = parser.parse_args()
 
@@ -301,6 +389,7 @@ def main():
     rolling_context = ""
     rolling_context_en = ""
     last_draft_text = ""
+    recent_en_lines = deque(maxlen=2)
 
     try:
         with sd.InputStream(samplerate=SAMPLERATE, channels=args.channels, callback=callback, blocksize=BLOCK_SIZE, device=device_index):
@@ -339,6 +428,7 @@ def main():
                         speech_started = False
                         rolling_context = ""
                         rolling_context_en = ""
+                        recent_en_lines.clear()
                         last_draft_text = ""
                         last_change_time = time.time()
                         if args.stream:
@@ -374,6 +464,22 @@ def main():
                             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:
@@ -385,6 +491,7 @@ def main():
                                 original_text = ""
                                 rolling_context = ""
                                 rolling_context_en = ""
+                                recent_en_lines.clear()
 
                             if original_text:
                                 print(f"\n[{detected_lang.upper()}]: {original_text}")
@@ -415,15 +522,35 @@ def main():
                                     english_text = bridge_result['text'].strip()
                                     if args.context and english_text:
                                         rolling_context_en = (rolling_context_en + " " + english_text)[-200:].strip()
-                                    if args.en:
-                                        payload["en"] = english_text
-                                        print(f"[EN]: {english_text}")
                                 else:
                                     english_text = original_text
                                     if args.context and english_text:
                                         rolling_context_en = (rolling_context_en + " " + english_text)[-200:].strip()
-                                    if args.en:
+
-                                        payload["en"] = english_text
+                                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, args.glossary_pair)
+
+                                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
+                                        english_for_caption = new_line
+                                        print(f"[EN-REV]: {old_line} -> {new_line}")
+                                    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:
@@ -431,8 +558,8 @@ def main():
                                     rolling_context = (rolling_context + " " + original_text)[-200:].strip()
 
                                 # 3. Translate from English to other languages
-                                if english_text and translation_engines:
+                                if english_for_translation and translation_engines:
-                                    english_chunks = split_text_for_translation(english_text, max_chars=args.mt_max_chars)
+                                    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')
@@ -508,6 +635,7 @@ def main():
                                     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")