faster-whisper-webui-v2 / src /prompts /jsonPromptStrategy.py
aadnk's picture
Use type hints instead
f641aea
import json
from typing import Dict
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] = dict()
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)