feat: add hallucination detection and watchdog for stuck transcriptions
This commit is contained in:
@@ -34,6 +34,10 @@ Run the script using `python3 transcribe.py` with optional flags.
|
||||
```bash
|
||||
python3 transcribe.py -l
|
||||
```
|
||||
- **Caption system audio (speakers) using a loopback device:**
|
||||
```bash
|
||||
python3 transcribe.py --loopback -es
|
||||
```
|
||||
- **Transcribe and translate to Spanish (screen only):**
|
||||
```bash
|
||||
python3 transcribe.py -es
|
||||
@@ -55,6 +59,7 @@ Run the script using `python3 transcribe.py` with optional flags.
|
||||
- `-i`, `--ingest`: Enable data transmission to the remote server.
|
||||
- `-l`, `--list-devices`: Show available audio devices and exit.
|
||||
- `-d`, `--device [ID]`: Input device index (bypasses selection prompt).
|
||||
- `--loopback`: Automatically select a loopback device (e.g., BlackHole, Stereo Mix) to caption system audio.
|
||||
- `-q`, `--quantize`: Use 4-bit quantized Whisper model for faster transcription (Strategy 3).
|
||||
- `-s`, `--stream`: Enable real-time streaming transcription/draft mode (Strategy 1).
|
||||
- `-c`, `--context`: Enable prompt caching/rolling context to help the model maintain sentence continuity across chunks.
|
||||
|
||||
+106
-3
@@ -1,10 +1,12 @@
|
||||
import sys
|
||||
import time
|
||||
import re
|
||||
import requests
|
||||
import threading
|
||||
import json
|
||||
import argparse
|
||||
from unittest.mock import MagicMock
|
||||
from collections import Counter
|
||||
|
||||
# Comprehensive workaround for missing _lzma in some Python builds
|
||||
try:
|
||||
@@ -45,11 +47,38 @@ CHANNELS = 1
|
||||
SAMPLERATE = 16000
|
||||
BLOCK_SIZE = 512
|
||||
VAD_THRESHOLD = 0.5
|
||||
BUFFER_LIMIT = SAMPLERATE * 30
|
||||
|
||||
audio_queue = queue.Queue()
|
||||
ingest_queue = queue.Queue()
|
||||
|
||||
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)
|
||||
@@ -91,11 +120,13 @@ def main():
|
||||
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)")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
@@ -107,6 +138,7 @@ def main():
|
||||
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}")
|
||||
|
||||
@@ -132,9 +164,28 @@ def main():
|
||||
print("Models loaded.")
|
||||
|
||||
# 2. Select Audio Device
|
||||
device_index = None
|
||||
if args.device is not None:
|
||||
device_index = args.device
|
||||
else:
|
||||
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:
|
||||
@@ -144,12 +195,19 @@ def main():
|
||||
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 = ""
|
||||
last_draft_text = ""
|
||||
|
||||
try:
|
||||
with sd.InputStream(samplerate=SAMPLERATE, channels=CHANNELS, callback=callback, blocksize=BLOCK_SIZE, device=device_index):
|
||||
@@ -161,6 +219,7 @@ def main():
|
||||
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,
|
||||
@@ -170,12 +229,28 @@ def main():
|
||||
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 = ""
|
||||
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:
|
||||
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:
|
||||
@@ -194,11 +269,18 @@ def main():
|
||||
original_text = transcription_result['text'].strip()
|
||||
detected_lang = transcription_result.get('language', args.lang if args.lang else 'en')
|
||||
|
||||
if is_hallucination(original_text):
|
||||
print(f"\n[SYSTEM]: Hallucination detected, ignoring and resetting context.")
|
||||
original_text = ""
|
||||
rolling_context = ""
|
||||
|
||||
if original_text:
|
||||
print(f"\n[{detected_lang.upper()}]: {original_text}")
|
||||
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):
|
||||
@@ -251,6 +333,8 @@ def main():
|
||||
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
|
||||
@@ -260,13 +344,32 @@ def main():
|
||||
|
||||
draft_result = mlx_whisper.transcribe(current_audio, **draft_kwargs)
|
||||
draft_text = draft_result['text'].strip()
|
||||
|
||||
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 = ""
|
||||
last_change_time = time.time()
|
||||
last_draft_text = ""
|
||||
sys.stdout.write("\r\033[K")
|
||||
sys.stdout.flush()
|
||||
continue
|
||||
|
||||
last_stream_time = time.time()
|
||||
|
||||
|
||||
Reference in New Issue
Block a user