Skip to content

feat: openvino support #106

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion pywhispercpp/constants.py
Original file line number Diff line number Diff line change
Expand Up @@ -139,7 +139,7 @@
},
'max_len': {
'type': int,
'description': "max segment length in characters",
'description': "max segment length in characters, note: token_timestamps needs to be set to True for this to work",
'options': None,
'default': 0
},
Expand Down
17 changes: 17 additions & 0 deletions pywhispercpp/model.py
Original file line number Diff line number Diff line change
Expand Up @@ -70,6 +70,10 @@ def __init__(self,
models_dir: str = None,
params_sampling_strategy: int = 0,
redirect_whispercpp_logs_to: Union[bool, TextIO, str, None] = False,
use_openvino: bool = False,
openvino_model_path: str = None,
openvino_device: str = 'CPU',
openvino_cache_dir: str = None,
**params):
"""
:param model: The name of the model, one of the [AVAILABLE_MODELS](/pywhispercpp/#pywhispercpp.constants.AVAILABLE_MODELS),
Expand All @@ -78,6 +82,10 @@ def __init__(self,
exist, default to [MODELS_DIR](/pywhispercpp/#pywhispercpp.constants.MODELS_DIR) <user_data_dir/pywhsipercpp/models>
:param params_sampling_strategy: 0 -> GREEDY, else BEAM_SEARCH
:param redirect_whispercpp_logs_to: where to redirect the whisper.cpp logs, default to False (no redirection), accepts str file path, sys.stdout, sys.stderr, or use None to redirect to devnull
:param use_openvino: whether to use OpenVINO or not
:param openvino_model_path: path to the OpenVINO model
:param openvino_device: OpenVINO device, default to CPU
:param openvino_cache_dir: OpenVINO cache directory
:param params: keyword arguments for different whisper.cpp parameters,
see [PARAMS_SCHEMA](/pywhispercpp/#pywhispercpp.constants.PARAMS_SCHEMA)
"""
Expand All @@ -90,8 +98,13 @@ def __init__(self,
pw.whisper_sampling_strategy.WHISPER_SAMPLING_BEAM_SEARCH
self._params = pw.whisper_full_default_params(self._sampling_strategy)
# assign params
self.params = params
self._set_params(params)
self.redirect_whispercpp_logs_to = redirect_whispercpp_logs_to
self.use_openvino = use_openvino
self.openvino_model_path = openvino_model_path
self.openvino_device = openvino_device
self.openvino_cache_dir = openvino_cache_dir
# init the model
self._init_model()

Expand Down Expand Up @@ -228,6 +241,10 @@ def _init_model(self) -> None:
logger.info("Initializing the model ...")
with utils.redirect_stderr(to=self.redirect_whispercpp_logs_to):
self._ctx = pw.whisper_init_from_file(self.model_path)
if self.use_openvino:
pw.whisper_ctx_init_openvino_encoder(self._ctx, self.openvino_model_path, self.openvino_device, self.openvino_cache_dir)



def _set_params(self, kwargs: dict) -> None:
"""
Expand Down
9 changes: 9 additions & 0 deletions src/main.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -276,6 +276,12 @@ float whisper_full_get_token_p_wrapper(struct whisper_context_wrapper * ctx, int
return whisper_full_get_token_p(ctx->ptr, i_segment, i_token);
}

int whisper_ctx_init_openvino_encoder_wrapper(struct whisper_context_wrapper * ctx, const char * model_path,
const char * device,
const char * cache_dir){
return whisper_ctx_init_openvino_encoder(ctx->ptr, model_path, device, cache_dir);
}

class WhisperFullParamsWrapper : public whisper_full_params {
std::string initial_prompt_str;
std::string suppress_regex_str;
Expand Down Expand Up @@ -656,6 +662,9 @@ PYBIND11_MODULE(_pywhispercpp, m) {

m.def("whisper_full_get_token_p", &whisper_full_get_token_p_wrapper, "Get the probability of the specified token in the specified segment.");

m.def("whisper_ctx_init_openvino_encoder", &whisper_ctx_init_openvino_encoder_wrapper, "Given a context, enable use of OpenVINO for encode inference.");


////////////////////////////////////////////////////////////////////////////

m.def("whisper_bench_memcpy", &whisper_bench_memcpy, "Temporary helpers needed for exposing ggml interface");
Expand Down
Loading