From 422bfab10452c093e8d63f1e3dce0a5603019c22 Mon Sep 17 00:00:00 2001 From: Sinan Date: Wed, 13 Nov 2024 15:26:00 +0100 Subject: [PATCH 1/2] Added encoder output --- faster_whisper/transcribe.py | 20 +++++++++++++++++--- 1 file changed, 17 insertions(+), 3 deletions(-) diff --git a/faster_whisper/transcribe.py b/faster_whisper/transcribe.py index 199bb099..58452415 100644 --- a/faster_whisper/transcribe.py +++ b/faster_whisper/transcribe.py @@ -9,7 +9,7 @@ from dataclasses import asdict, dataclass from inspect import signature from math import ceil -from typing import BinaryIO, Iterable, List, Optional, Tuple, Union +from typing import BinaryIO, Iterable, List, Optional, Tuple, Union, Any from warnings import warn import ctranslate2 @@ -114,6 +114,20 @@ class TranscriptionInfo: transcription_options: TranscriptionOptions vad_options: VadOptions +@dataclass +class TranscriptionResult: + segments: Iterable["Segment"] + info: "TranscriptionInfo" + encoder_output: Any + + def __iter__(self) -> Iterable: + # Only include the original values for backwards compatibility + yield self.segments + yield self.info + + def all(self) -> Tuple[Iterable["Segment"], "TranscriptionInfo", Any]: + # Access all three components + return self.segments, self.info, self.encoder_output # The code below is originally from HF pipeline and is used in whisper-x # (https://github.com/m-bain/whisperX) and adapted for faster_whisper @@ -694,7 +708,7 @@ def transcribe( hotwords: Optional[str] = None, language_detection_threshold: Optional[float] = None, language_detection_segments: int = 1, - ) -> Tuple[Iterable[Segment], TranscriptionInfo]: + ) -> TranscriptionResult: """Transcribes an input file. Arguments: @@ -973,7 +987,7 @@ def transcribe( vad_options=vad_parameters, all_language_probs=all_language_probs, ) - return segments, info + return TranscriptionResult(segments=segments, info=info, encoder_output=encoder_output) def _split_segments_by_timestamps( self, From b8839a6eb3fb796c2fd85faccbe78c4ac079e9c0 Mon Sep 17 00:00:00 2001 From: Sinan Date: Wed, 13 Nov 2024 15:43:51 +0100 Subject: [PATCH 2/2] blank line --- faster_whisper/transcribe.py | 1 + 1 file changed, 1 insertion(+) diff --git a/faster_whisper/transcribe.py b/faster_whisper/transcribe.py index 58452415..4d337339 100644 --- a/faster_whisper/transcribe.py +++ b/faster_whisper/transcribe.py @@ -114,6 +114,7 @@ class TranscriptionInfo: transcription_options: TranscriptionOptions vad_options: VadOptions + @dataclass class TranscriptionResult: segments: Iterable["Segment"]