Adding JSON initial prompt

By selecting "json_prompt_mode", you can
customize the prompt to each segment.

For instance:
[
{"segment_index": 0, "prompt": "Hello, how are you?"},
{"segment_index": 1, "prompt": "I'm doing well, how are you?"},
{"segment_index": 2, "prompt": "{0} Fine, thank you.", "format_prompt": true}
]

app.py

@@ -13,12 +13,14 @@ import numpy as np
 
 import torch
 
-from src.config import ApplicationConfig, VadInitialPromptMode
+from src.config import VAD_INITIAL_PROMPT_MODE_VALUES, ApplicationConfig, VadInitialPromptMode
 from src.hooks.progressListener import ProgressListener
 from src.hooks.subTaskProgressListener import SubTaskProgressListener
 from src.hooks.whisperProgressHook import create_progress_listener_handle
 from src.languages import get_language_names
 from src.modelCache import ModelCache
+from src.prompts.jsonPromptStrategy import JsonPromptStrategy
+from src.prompts.prependPromptStrategy import PrependPromptStrategy
 from src.source import get_audio_source_collection
 from src.vadParallel import ParallelContext, ParallelTranscription
 
@@ -271,8 +273,18 @@ class WhisperTranscriber:
         if ('task' in decodeOptions):
             task = decodeOptions.pop('task')
 
+        if (vadOptions.vadInitialPromptMode == VadInitialPromptMode.PREPEND_ALL_SEGMENTS or 
+            vadOptions.vadInitialPromptMode == VadInitialPromptMode.PREPREND_FIRST_SEGMENT):
+            # Prepend initial prompt
+            prompt_strategy = PrependPromptStrategy(initial_prompt, vadOptions.vadInitialPromptMode)
+        elif (vadOptions.vadInitialPromptMode == VadInitialPromptMode.JSON_PROMPT_MODE):
+            # Use a JSON format to specify the prompt for each segment
+            prompt_strategy = JsonPromptStrategy(initial_prompt)
+        else:
+            raise ValueError("Invalid vadInitialPromptMode: " + vadOptions.vadInitialPromptMode)
+
         # Callable for processing an audio file
-        whisperCallable = model.create_callback(language, task, initial_prompt, initial_prompt_mode=vadOptions.vadInitialPromptMode, **decodeOptions)
+        whisperCallable = model.create_callback(language, task, prompt_strategy=prompt_strategy, **decodeOptions)
 
         # The results
         if (vadOptions.vad == 'silero-vad'):
@@ -519,7 +531,7 @@ def create_ui(app_config: ApplicationConfig):
         *common_vad_inputs(),
         gr.Number(label="VAD - Padding (s)", precision=None, value=app_config.vad_padding),
         gr.Number(label="VAD - Prompt Window (s)", precision=None, value=app_config.vad_prompt_window),
-        gr.Dropdown(choices=["prepend_first_segment", "prepend_all_segments"], value=app_config.vad_initial_prompt_mode, label="VAD - Initial Prompt Mode"),
+        gr.Dropdown(choices=VAD_INITIAL_PROMPT_MODE_VALUES, label="VAD - Initial Prompt Mode"),
         
         *common_word_timestamps_inputs(),
         gr.Text(label="Word Timestamps - Prepend Punctuations", value=app_config.prepend_punctuations),
@@ -580,7 +592,7 @@ if __name__ == '__main__':
                         help="The default model name.") # medium
     parser.add_argument("--default_vad", type=str, default=default_app_config.default_vad, \
                         help="The default VAD.") # silero-vad
-    parser.add_argument("--vad_initial_prompt_mode", type=str, default=default_app_config.vad_initial_prompt_mode, choices=["prepend_all_segments", "prepend_first_segment"], \
+    parser.add_argument("--vad_initial_prompt_mode", type=str, default=default_app_config.vad_initial_prompt_mode, choices=VAD_INITIAL_PROMPT_MODE_VALUES, \
                         help="Whether or not to prepend the initial prompt to each VAD segment (prepend_all_segments), or just the first segment (prepend_first_segment)") # prepend_first_segment
     parser.add_argument("--vad_parallel_devices", type=str, default=default_app_config.vad_parallel_devices, \
                         help="A commma delimited list of CUDA devices to use for parallel processing. If None, disable parallel processing.") # ""

cli.py

@@ -7,7 +7,7 @@ import numpy as np
 
 import torch
 from app import VadOptions, WhisperTranscriber
-from src.config import ApplicationConfig, VadInitialPromptMode
+from src.config import VAD_INITIAL_PROMPT_MODE_VALUES, ApplicationConfig, VadInitialPromptMode
 from src.download import download_url
 from src.languages import get_language_names
 
@@ -47,7 +47,7 @@ def cli():
 
     parser.add_argument("--vad", type=str, default=app_config.default_vad, choices=["none", "silero-vad", "silero-vad-skip-gaps", "silero-vad-expand-into-gaps", "periodic-vad"], \
                         help="The voice activity detection algorithm to use") # silero-vad
-    parser.add_argument("--vad_initial_prompt_mode", type=str, default=app_config.vad_initial_prompt_mode, choices=["prepend_all_segments", "prepend_first_segment"], \
+    parser.add_argument("--vad_initial_prompt_mode", type=str, default=app_config.vad_initial_prompt_mode, choices=VAD_INITIAL_PROMPT_MODE_VALUES, \
                         help="Whether or not to prepend the initial prompt to each VAD segment (prepend_all_segments), or just the first segment (prepend_first_segment)") # prepend_first_segment
     parser.add_argument("--vad_merge_window", type=optional_float, default=app_config.vad_merge_window, \
                         help="The window size (in seconds) to merge voice segments")

src/config.py

@@ -24,9 +24,12 @@ class ModelConfig:
         self.path = path
         self.type = type
 
+VAD_INITIAL_PROMPT_MODE_VALUES=["prepend_all_segments", "prepend_first_segment", "json_prompt_mode"]
+
 class VadInitialPromptMode(Enum):
     PREPEND_ALL_SEGMENTS = 1
     PREPREND_FIRST_SEGMENT = 2
+    JSON_PROMPT_MODE = 3
 
     @staticmethod
     def from_string(s: str):
@@ -36,6 +39,8 @@ class VadInitialPromptMode(Enum):
             return VadInitialPromptMode.PREPEND_ALL_SEGMENTS
         elif normalized == "prepend_first_segment":
             return VadInitialPromptMode.PREPREND_FIRST_SEGMENT
+        elif normalized == "json_prompt_mode":
+            return VadInitialPromptMode.JSON_PROMPT_MODE
         else:
             raise ValueError(f"Invalid value for VadInitialPromptMode: {s}")
 

src/prompts/abstractPromptStrategy.py

@@ -0,0 +1,73 @@
+import abc
+
+
+class AbstractPromptStrategy:
+    """
+    Represents a strategy for generating prompts for a given audio segment.
+
+    Note that the strategy must be picklable, as it will be serialized and sent to the workers.
+    """
+    
+    @abc.abstractmethod
+    def get_segment_prompt(self, segment_index: int, whisper_prompt: str, detected_language: str) -> str:
+        """
+        Retrieves the prompt for a given segment.
+
+        Parameters
+        ----------
+        segment_index: int
+            The index of the segment.
+        whisper_prompt: str
+            The prompt for the segment generated by Whisper. This is typically concatenated with the initial prompt.
+        detected_language: str
+            The language detected for the segment.
+        """
+        pass
+
+    @abc.abstractmethod
+    def on_segment_finished(self, segment_index: int, whisper_prompt: str, detected_language: str, result: dict):
+        """
+        Called when a segment has finished processing.
+        
+        Parameters
+        ----------
+        segment_index: int
+            The index of the segment.
+        whisper_prompt: str
+            The prompt for the segment generated by Whisper. This is typically concatenated with the initial prompt.
+        detected_language: str
+            The language detected for the segment.
+        result: dict
+            The result of the segment. It has the following format:
+                {
+                    "text": str,
+                    "segments": [
+                        {
+                            "text": str,
+                            "start": float,
+                            "end": float,
+                            "words": [words],
+                        }
+                    ],
+                    "language": str,
+                }
+        """
+        pass
+
+    def _concat_prompt(self, prompt1, prompt2):
+        """
+        Concatenates two prompts.
+
+        Parameters
+        ----------
+        prompt1: str
+            The first prompt.
+        prompt2: str
+            The second prompt.
+        """
+        if (prompt1 is None):
+            return prompt2
+        elif (prompt2 is None):
+            return prompt1
+        else:
+            return prompt1 + " " + prompt2

src/prompts/jsonPromptStrategy.py

@@ -0,0 +1,48 @@
+import json
+from src.prompts.abstractPromptStrategy import AbstractPromptStrategy
+
+
+class JsonPromptSegment():
+    def __init__(self, segment_index: int, prompt: str, format_prompt: bool = False):
+        self.prompt = prompt
+        self.segment_index = segment_index
+        self.format_prompt = format_prompt
+
+class JsonPromptStrategy(AbstractPromptStrategy):
+    def __init__(self, initial_json_prompt: str):
+        """
+        Parameters
+        ----------
+            initial_json_prompt: str
+                The initial prompts for each segment in JSON form.
+
+                Format:
+                [
+                    {"segment_index": 0, "prompt": "Hello, how are you?"},
+                    {"segment_index": 1, "prompt": "I'm doing well, how are you?"},
+                    {"segment_index": 2, "prompt": "{0} Fine, thank you.", "format_prompt": true}
+                ]
+                    
+        """
+        parsed_json = json.loads(initial_json_prompt)
+        self.segment_lookup = dict[str, JsonPromptSegment]()
+        
+        for prompt_entry in parsed_json:
+            segment_index = prompt_entry["segment_index"]
+            prompt = prompt_entry["prompt"]
+            format_prompt = prompt_entry.get("format_prompt", False)
+            self.segment_lookup[str(segment_index)] = JsonPromptSegment(segment_index, prompt, format_prompt)
+
+    def get_segment_prompt(self, segment_index: int, whisper_prompt: str, detected_language: str) -> str:
+        # Lookup prompt
+        prompt = self.segment_lookup.get(str(segment_index), None)
+
+        if (prompt is None):
+            # No prompt found, return whisper prompt
+            print(f"Could not find prompt for segment {segment_index}, returning whisper prompt")
+            return whisper_prompt
+
+        if (prompt.format_prompt):
+            return prompt.prompt.format(whisper_prompt)
+        else:
+            return self._concat_prompt(prompt.prompt, whisper_prompt)

src/prompts/prependPromptStrategy.py

@@ -0,0 +1,31 @@
+from src.config import VadInitialPromptMode
+from src.prompts.abstractPromptStrategy import AbstractPromptStrategy
+
+class PrependPromptStrategy(AbstractPromptStrategy):
+    """
+    A simple prompt strategy that prepends a single prompt to all segments of audio, or prepends the prompt to the first segment of audio.
+    """
+    def __init__(self, initial_prompt: str, initial_prompt_mode: VadInitialPromptMode):
+        """
+        Parameters
+        ----------
+            initial_prompt: str
+                The initial prompt to use for the transcription.
+            initial_prompt_mode: VadInitialPromptMode
+                The mode to use for the initial prompt. If set to PREPEND_FIRST_SEGMENT, the initial prompt will be prepended to the first segment of audio.
+                If set to PREPEND_ALL_SEGMENTS, the initial prompt will be prepended to all segments of audio.
+        """
+        self.initial_prompt = initial_prompt
+        self.initial_prompt_mode = initial_prompt_mode
+
+        # This is a simple prompt strategy, so we only support these two modes
+        if initial_prompt_mode not in [VadInitialPromptMode.PREPEND_ALL_SEGMENTS, VadInitialPromptMode.PREPREND_FIRST_SEGMENT]:
+            raise ValueError(f"Unsupported initial prompt mode {initial_prompt_mode}")
+
+    def get_segment_prompt(self, segment_index: int, whisper_prompt: str, detected_language: str) -> str:
+        if (self.initial_prompt_mode == VadInitialPromptMode.PREPEND_ALL_SEGMENTS):
+            return self._concat_prompt(self.initial_prompt, whisper_prompt)
+        elif (self.initial_prompt_mode == VadInitialPromptMode.PREPREND_FIRST_SEGMENT):
+            return self._concat_prompt(self.initial_prompt, whisper_prompt) if segment_index == 0 else whisper_prompt
+        else:
+            raise ValueError(f"Unknown initial prompt mode {self.initial_prompt_mode}")

src/whisper/abstractWhisperContainer.py

@@ -1,11 +1,16 @@
 import abc
 from typing import List
+
 from src.config import ModelConfig, VadInitialPromptMode
 
 from src.hooks.progressListener import ProgressListener
 from src.modelCache import GLOBAL_MODEL_CACHE, ModelCache
+from src.prompts.abstractPromptStrategy import AbstractPromptStrategy
 
 class AbstractWhisperCallback:
+    def __init__(self):
+        self.__prompt_mode_gpt = None
+
     @abc.abstractmethod
     def invoke(self, audio, segment_index: int, prompt: str, detected_language: str, progress_listener: ProgressListener = None):
         """
@@ -24,23 +29,6 @@ class AbstractWhisperCallback:
         """
         raise NotImplementedError()
 
-    def _get_initial_prompt(self, initial_prompt: str, initial_prompt_mode: VadInitialPromptMode, 
-                               prompt: str, segment_index: int):
-        if (initial_prompt_mode == VadInitialPromptMode.PREPEND_ALL_SEGMENTS):
-            return self._concat_prompt(initial_prompt, prompt)
-        elif (initial_prompt_mode == VadInitialPromptMode.PREPREND_FIRST_SEGMENT):
-            return self._concat_prompt(initial_prompt, prompt) if segment_index == 0 else prompt
-        else:
-            raise ValueError(f"Unknown initial prompt mode {initial_prompt_mode}")
-
-    def _concat_prompt(self, prompt1, prompt2):
-        if (prompt1 is None):
-            return prompt2
-        elif (prompt2 is None):
-            return prompt1
-        else:
-            return prompt1 + " " + prompt2
-
 class AbstractWhisperContainer:
     def __init__(self, model_name: str, device: str = None, compute_type: str = "float16",
                  download_root: str = None,
@@ -75,8 +63,8 @@ class AbstractWhisperContainer:
         pass
 
     @abc.abstractmethod
-    def create_callback(self, language: str = None, task: str = None, initial_prompt: str = None, 
-                        initial_prompt_mode: VadInitialPromptMode = VadInitialPromptMode.PREPREND_FIRST_SEGMENT, 
+    def create_callback(self, language: str = None, task: str = None, 
+                        prompt_strategy: AbstractPromptStrategy = None, 
                         **decodeOptions: dict) -> AbstractWhisperCallback:
         """
         Create a WhisperCallback object that can be used to transcript audio files.
@@ -87,11 +75,8 @@ class AbstractWhisperContainer:
             The target language of the transcription. If not specified, the language will be inferred from the audio content.
         task: str
             The task - either translate or transcribe.
-        initial_prompt: str
-            The initial prompt to use for the transcription.
-        initial_prompt_mode: VadInitialPromptMode
-            The mode to use for the initial prompt. If set to PREPEND_FIRST_SEGMENT, the initial prompt will be prepended to the first segment of audio.
-            If set to PREPEND_ALL_SEGMENTS, the initial prompt will be prepended to all segments of audio.
+        prompt_strategy: AbstractPromptStrategy
+            The prompt strategy to use for the transcription.
         decodeOptions: dict
             Additional options to pass to the decoder. Must be pickleable.
 

src/whisper/fasterWhisperContainer.py

@@ -6,6 +6,7 @@ from src.config import ModelConfig, VadInitialPromptMode
 from src.hooks.progressListener import ProgressListener
 from src.languages import get_language_from_name
 from src.modelCache import ModelCache
+from src.prompts.abstractPromptStrategy import AbstractPromptStrategy
 from src.whisper.abstractWhisperContainer import AbstractWhisperCallback, AbstractWhisperContainer
 from src.utils import format_timestamp
 
@@ -56,8 +57,8 @@ class FasterWhisperContainer(AbstractWhisperContainer):
         model = WhisperModel(model_url, device=device, compute_type=self.compute_type)
         return model
 
-    def create_callback(self, language: str = None, task: str = None, initial_prompt: str = None, 
-                        initial_prompt_mode: VadInitialPromptMode = VadInitialPromptMode.PREPREND_FIRST_SEGMENT, 
+    def create_callback(self, language: str = None, task: str = None, 
+                        prompt_strategy: AbstractPromptStrategy = None, 
                         **decodeOptions: dict) -> AbstractWhisperCallback:
         """
         Create a WhisperCallback object that can be used to transcript audio files.
@@ -68,11 +69,8 @@ class FasterWhisperContainer(AbstractWhisperContainer):
             The target language of the transcription. If not specified, the language will be inferred from the audio content.
         task: str
             The task - either translate or transcribe.
-        initial_prompt: str
-            The initial prompt to use for the transcription.
-        initial_prompt_mode: VadInitialPromptMode
-            The mode to use for the initial prompt. If set to PREPEND_FIRST_SEGMENT, the initial prompt will be prepended to the first segment of audio.
-            If set to PREPEND_ALL_SEGMENTS, the initial prompt will be prepended to all segments of audio.
+        prompt_strategy: AbstractPromptStrategy
+            The prompt strategy to use. If not specified, the prompt from Whisper will be used.
         decodeOptions: dict
             Additional options to pass to the decoder. Must be pickleable.
 
@@ -80,17 +78,16 @@ class FasterWhisperContainer(AbstractWhisperContainer):
         -------
         A WhisperCallback object.
         """
-        return FasterWhisperCallback(self, language=language, task=task, initial_prompt=initial_prompt, initial_prompt_mode=initial_prompt_mode, **decodeOptions)
+        return FasterWhisperCallback(self, language=language, task=task, prompt_strategy=prompt_strategy, **decodeOptions)
 
 class FasterWhisperCallback(AbstractWhisperCallback):
     def __init__(self, model_container: FasterWhisperContainer, language: str = None, task: str = None, 
-                 initial_prompt: str = None, initial_prompt_mode: VadInitialPromptMode=VadInitialPromptMode.PREPREND_FIRST_SEGMENT, 
+                 prompt_strategy: AbstractPromptStrategy = None, 
                  **decodeOptions: dict):
         self.model_container = model_container
         self.language = language
         self.task = task
-        self.initial_prompt = initial_prompt
-        self.initial_prompt_mode = initial_prompt_mode
+        self.prompt_strategy = prompt_strategy
         self.decodeOptions = decodeOptions
 
         self._printed_warning = False
@@ -138,7 +135,8 @@ class FasterWhisperCallback(AbstractWhisperCallback):
         # See if supress_tokens is a string - if so, convert it to a list of ints
         decodeOptions["suppress_tokens"] = self._split_suppress_tokens(suppress_tokens)
 
-        initial_prompt = self._get_initial_prompt(self.initial_prompt, self.initial_prompt_mode, prompt, segment_index)
+        initial_prompt = self.prompt_strategy.get_segment_prompt(segment_index, prompt, detected_language) \
+                           if self.prompt_strategy else prompt
 
         segments_generator, info = model.transcribe(audio, \
             language=language_code if language_code else detected_language, task=self.task, \
@@ -184,6 +182,10 @@ class FasterWhisperCallback(AbstractWhisperCallback):
             "duration": info.duration if info else None
         }
 
+        # If we have a prompt strategy, we need to increment the current prompt
+        if self.prompt_strategy:
+            self.prompt_strategy.on_segment_finished(segment_index, prompt, detected_language, result)
+
         if progress_listener is not None:
             progress_listener.on_finished()
         return result

src/whisper/whisperContainer.py

@@ -15,6 +15,7 @@ from src.config import ModelConfig, VadInitialPromptMode
 from src.hooks.whisperProgressHook import create_progress_listener_handle
 
 from src.modelCache import GLOBAL_MODEL_CACHE, ModelCache
+from src.prompts.abstractPromptStrategy import AbstractPromptStrategy
 from src.utils import download_file
 from src.whisper.abstractWhisperContainer import AbstractWhisperCallback, AbstractWhisperContainer
 
@@ -69,8 +70,8 @@ class WhisperContainer(AbstractWhisperContainer):
 
         return whisper.load_model(model_path, device=self.device, download_root=self.download_root)
 
-    def create_callback(self, language: str = None, task: str = None, initial_prompt: str = None, 
-                        initial_prompt_mode: VadInitialPromptMode = VadInitialPromptMode.PREPREND_FIRST_SEGMENT, 
+    def create_callback(self, language: str = None, task: str = None, 
+                        prompt_strategy: AbstractPromptStrategy = None,
                         **decodeOptions: dict) -> AbstractWhisperCallback:
         """
         Create a WhisperCallback object that can be used to transcript audio files.
@@ -81,11 +82,8 @@ class WhisperContainer(AbstractWhisperContainer):
             The target language of the transcription. If not specified, the language will be inferred from the audio content.
         task: str
             The task - either translate or transcribe.
-        initial_prompt: str
-            The initial prompt to use for the transcription.
-        initial_prompt_mode: VadInitialPromptMode
-            The mode to use for the initial prompt. If set to PREPEND_FIRST_SEGMENT, the initial prompt will be prepended to the first segment of audio.
-            If set to PREPEND_ALL_SEGMENTS, the initial prompt will be prepended to all segments of audio.
+        prompt_strategy: AbstractPromptStrategy
+            The prompt strategy to use. If not specified, the prompt from Whisper will be used.
         decodeOptions: dict
             Additional options to pass to the decoder. Must be pickleable.
 
@@ -93,7 +91,7 @@ class WhisperContainer(AbstractWhisperContainer):
         -------
         A WhisperCallback object.
         """
-        return WhisperCallback(self, language=language, task=task, initial_prompt=initial_prompt, initial_prompt_mode=initial_prompt_mode, **decodeOptions)
+        return WhisperCallback(self, language=language, task=task, prompt_strategy=prompt_strategy, **decodeOptions)
 
     def _get_model_path(self, model_config: ModelConfig, root_dir: str = None):
         from src.conversion.hf_converter import convert_hf_whisper
@@ -162,13 +160,14 @@ class WhisperContainer(AbstractWhisperContainer):
         return model_config.path
 
 class WhisperCallback(AbstractWhisperCallback):
-    def __init__(self, model_container: WhisperContainer, language: str = None, task: str = None, initial_prompt: str = None, 
-                 initial_prompt_mode: VadInitialPromptMode=VadInitialPromptMode.PREPREND_FIRST_SEGMENT, **decodeOptions: dict):
+    def __init__(self, model_container: WhisperContainer, language: str = None, task: str = None, 
+                 prompt_strategy: AbstractPromptStrategy = None, 
+                 **decodeOptions: dict):
         self.model_container = model_container
         self.language = language
         self.task = task
-        self.initial_prompt = initial_prompt
-        self.initial_prompt_mode = initial_prompt_mode
+        self.prompt_strategy = prompt_strategy
+
         self.decodeOptions = decodeOptions
         
     def invoke(self, audio, segment_index: int, prompt: str, detected_language: str, progress_listener: ProgressListener = None):
@@ -201,11 +200,17 @@ class WhisperCallback(AbstractWhisperCallback):
         if self.model_container.compute_type in ["fp16", "float16"]:
             decodeOptions["fp16"] = True
 
-        initial_prompt = self._get_initial_prompt(self.initial_prompt, self.initial_prompt_mode, prompt, segment_index)
+        initial_prompt = self.prompt_strategy.get_segment_prompt(segment_index, prompt, detected_language) \
+                           if self.prompt_strategy else prompt
 
         result = model.transcribe(audio, \
             language=self.language if self.language else detected_language, task=self.task, \
             initial_prompt=initial_prompt, \
             **decodeOptions
         )
+
+        # If we have a prompt strategy, we need to increment the current prompt
+        if self.prompt_strategy:
+            self.prompt_strategy.on_segment_finished(segment_index, prompt, detected_language, result)
+
         return result