diff --git a/pywhispercpp/constants.py b/pywhispercpp/constants.py index bfe582f..f92aa76 100644 --- a/pywhispercpp/constants.py +++ b/pywhispercpp/constants.py @@ -302,12 +302,6 @@ 'options': None, 'default': {"beam_size": -1, "patience": -1.0} }, - 'extract_probability': { - 'type': bool, - 'description': 'calculate the geometric mean of token probabilities for each segment.', - 'options': None, - 'default': False - }, 'vad': { 'type': bool, 'description': 'Enable VAD', diff --git a/pywhispercpp/model.py b/pywhispercpp/model.py index 453e152..08be6bd 100644 --- a/pywhispercpp/model.py +++ b/pywhispercpp/model.py @@ -181,6 +181,7 @@ def transcribe(self, n_processors: Optional[int] = None, new_segment_callback: Optional[Callable[[Segment], None]] = None, abort_callback: Optional[Callable[[], bool]] = None, + extract_probability: bool = False, **params) -> List[Segment]: """ Transcribes the media provided as input and returns list of `Segment` objects. @@ -205,9 +206,6 @@ def transcribe(self, raise FileNotFoundError(media) audio = self._load_audio(media) - # Handle extract_probability parameter - self.extract_probability = params.pop('extract_probability', False) - # update params if any self._set_params(params) @@ -224,7 +222,7 @@ def transcribe(self, # run inference start_time = time() logger.info("Transcribing ...") - res = self._transcribe(audio, n_processors=n_processors) + res = self._transcribe(audio, n_processors=n_processors, extract_probability=extract_probability) end_time = time() logger.info(f"Inference time: {end_time - start_time:.3f} s") return res @@ -402,12 +400,14 @@ def _set_params(self, kwargs: dict) -> None: for param, value in normalized.items(): setattr(self._params, param, value) - def _transcribe(self, audio: np.ndarray, n_processors: Optional[int] = None): + def _transcribe(self, audio: np.ndarray, n_processors: Optional[int] = None, extract_probability: bool = False): """ Private method to call the whisper.cpp/whisper_full function :param audio: numpy array of audio data :param n_processors: if not None, it will run whisper.cpp/whisper_full_parallel with n_processors + :param extract_probability: If True, calculates the geometric mean of token probabilities for each segment, + providing a confidence score interpretable as a probability in [0, 1]. :return: """ @@ -416,7 +416,7 @@ def _transcribe(self, audio: np.ndarray, n_processors: Optional[int] = None): else: pw.whisper_full(self._ctx, self._params, audio, audio.size) n = pw.whisper_full_n_segments(self._ctx) - res = Model._get_segments(self._ctx, 0, n, self.extract_probability) + res = Model._get_segments(self._ctx, 0, n, extract_probability) return res @@ -528,4 +528,4 @@ def __del__(self): :return: None """ if self._ctx is not None: - pw.whisper_free(self._ctx) \ No newline at end of file + pw.whisper_free(self._ctx)