feat: set initial_prompt and vad_parameters in the first message

This commit is contained in:
Chen Hua
2024-01-08 15:33:55 +08:00
parent 7b2f5cff72
commit 72ead71eeb
+10 -4
View File
@@ -102,7 +102,9 @@ class TranscriptionServer:
language=options["language"], language=options["language"],
task=options["task"], task=options["task"],
client_uid=options["uid"], client_uid=options["uid"],
model_size=options["model_size"] model_size=options["model_size"],
initial_prompt=options["initial_prompt"],
vad_parameters=options["vad_parameters"]
) )
self.clients[websocket] = client self.clients[websocket] = client
@@ -190,7 +192,9 @@ class ServeClient:
multilingual=False, multilingual=False,
language=None, language=None,
client_uid=None, client_uid=None,
model_size="small" model_size="small",
initial_prompt=None,
vad_parameters=None
): ):
""" """
Initialize a ServeClient instance. Initialize a ServeClient instance.
@@ -218,6 +222,8 @@ class ServeClient:
self.language = language if self.multilingual else "en" self.language = language if self.multilingual else "en"
self.task = task self.task = task
self.websocket = websocket self.websocket = websocket
self.initial_prompt = initial_prompt
self.vad_parameters = vad_parameters or {"threshold": 0.5}
device = "cuda" if torch.cuda.is_available() else "cpu" device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -352,11 +358,11 @@ class ServeClient:
# whisper transcribe with prompt # whisper transcribe with prompt
result, info = self.transcriber.transcribe( result, info = self.transcriber.transcribe(
input_sample, input_sample,
initial_prompt=None, initial_prompt=self.initial_prompt,
language=self.language, language=self.language,
task=self.task, task=self.task,
vad_filter=True, vad_filter=True,
vad_parameters={"threshold": 0.5} vad_parameters=self.vad_parameters
) )
if self.language is None: if self.language is None: