feat: implement prompt caching (-c) and source language bypass (--lang)

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.
 
 ---
 

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}")