Adolfo Reyna
873 lines
43 KiB
Python
873 lines
43 KiB
Python
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()
|