feat: add intelligent english caption correction v1
This commit is contained in:
Adolfo Reyna
@@ -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.
|
||||
|
||||
+136
-8
@@ -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:
|
||||
english_chunks = split_text_for_translation(english_text, max_chars=args.mt_max_chars)
|
||||
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')
|
||||
@@ -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")
|
||||
|
||||
Reference in New Issue
Block a user