From 02793a93f80d73baeff95897bcfc533b47545bde Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=8D=8E=E6=99=A8?= Date: Mon, 1 Jan 2024 20:22:37 +0800 Subject: [PATCH] feat: Update transcriber to support large-v3 model with 128 mel filters --- whisper_live/transcriber.py | 21 ++++++++++++++++++++- 1 file changed, 20 insertions(+), 1 deletion(-) diff --git a/whisper_live/transcriber.py b/whisper_live/transcriber.py index 918dc2a..cf301d3 100644 --- a/whisper_live/transcriber.py +++ b/whisper_live/transcriber.py @@ -4,6 +4,8 @@ import itertools import logging import os import zlib +import json +from inspect import signature from typing import BinaryIO, Iterable, List, NamedTuple, Optional, Tuple, Union @@ -144,7 +146,8 @@ class WhisperModel: "openai/whisper-tiny" + ("" if self.model.is_multilingual else ".en") ) - self.feature_extractor = FeatureExtractor() + self.feat_kwargs = self._get_feature_kwargs(model_path) + self.feature_extractor = FeatureExtractor(**self.feat_kwargs) self.num_samples_per_token = self.feature_extractor.hop_length * 2 self.frames_per_second = ( self.feature_extractor.sampling_rate // self.feature_extractor.hop_length @@ -161,6 +164,22 @@ class WhisperModel: """The languages supported by the model.""" return list(_LANGUAGE_CODES) if self.model.is_multilingual else ["en"] + def _get_feature_kwargs(self, model_path) -> dict: + preprocessor_config_file = os.path.join(model_path, "preprocessor_config.json") + config = {} + if os.path.isfile(preprocessor_config_file): + try: + with open(preprocessor_config_file, "r", encoding="utf-8") as json_file: + config = json.load(json_file) + valid_keys = signature(FeatureExtractor.__init__).parameters.keys() + config = {k: v for k, v in config.items() if k in valid_keys} + except json.JSONDecodeError as e: + self.logger.warning( + "Could not load preprocessor_config.json: %s", str(e) + ) + + return config + def transcribe( self, audio: Union[str, BinaryIO, np.ndarray],