diff --git a/README.md b/README.md index 2cd2512..aa9a0cf 100644 --- a/README.md +++ b/README.md @@ -57,6 +57,8 @@ Run the script using `python3 transcribe.py` with optional flags. - `-d`, `--device [ID]`: Input device index (bypasses selection prompt). - `-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. +- `--lang [CODE]`: Hardcode the source language (e.g., `en`, `es`) to bypass automatic language detection for faster processing. --- diff --git a/transcribe.py b/transcribe.py index d42876d..5b916db 100644 --- a/transcribe.py +++ b/transcribe.py @@ -94,6 +94,8 @@ def main(): parser.add_argument("-d", "--device", type=int, help="Input device index") 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") args = parser.parse_args() @@ -147,6 +149,7 @@ def main(): audio_buffer = [] speech_started = False last_stream_time = time.time() + rolling_context = "" try: with sd.InputStream(samplerate=SAMPLERATE, channels=CHANNELS, callback=callback, blocksize=BLOCK_SIZE, device=device_index): @@ -179,10 +182,17 @@ def main(): 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 + # 1. Transcribe & Detect Language - transcription_result = mlx_whisper.transcribe(current_audio, path_or_hf_repo=WHISPER_MODEL) + transcription_result = mlx_whisper.transcribe(current_audio, **transcribe_kwargs) original_text = transcription_result['text'].strip() - detected_lang = transcription_result.get('language', 'en') + detected_lang = transcription_result.get('language', args.lang if args.lang else 'en') if original_text: print(f"\n[{detected_lang.upper()}]: {original_text}") @@ -196,7 +206,10 @@ def main(): # 2. Bridge to English if not already English if detected_lang != "en": - bridge_result = mlx_whisper.transcribe(current_audio, path_or_hf_repo=WHISPER_MODEL, task="translate") + bridge_kwargs = {"path_or_hf_repo": WHISPER_MODEL, "task": "translate"} + if args.lang: + bridge_kwargs["language"] = args.lang + bridge_result = mlx_whisper.transcribe(current_audio, **bridge_kwargs) english_text = bridge_result['text'].strip() if args.en: payload["en"] = english_text @@ -206,6 +219,11 @@ def main(): if args.en: payload["en"] = english_text + # 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_text and translation_engines: # Limit input length @@ -236,7 +254,11 @@ def main(): elif args.stream and (time.time() - last_stream_time) > 1.0: # Draft transcription - draft_result = mlx_whisper.transcribe(current_audio, path_or_hf_repo=WHISPER_MODEL) + 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 + + draft_result = mlx_whisper.transcribe(current_audio, **draft_kwargs) draft_text = draft_result['text'].strip() if draft_text: sys.stdout.write(f"\r\033[K[DRAFT]: {draft_text}")