From fc0f40981ded1c74db4ce57878b4d7f56fec1b9c Mon Sep 17 00:00:00 2001 From: EliEron Date: Mon, 26 Sep 2022 13:52:28 +0200 Subject: [PATCH] Write each sentence as a separate line for the txt output (#101) * Write each sentence as a separate line for the txt output Write each sentence as a separate line for the txt output * Update utils.py Co-authored-by: EliEron Co-authored-by: Jong Wook Kim --- whisper/transcribe.py | 4 ++-- whisper/utils.py | 5 +++++ 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/whisper/transcribe.py b/whisper/transcribe.py index 82a705649..8f34b0ef3 100644 --- a/whisper/transcribe.py +++ b/whisper/transcribe.py @@ -10,7 +10,7 @@ from .audio import SAMPLE_RATE, N_FRAMES, HOP_LENGTH, pad_or_trim, log_mel_spectrogram from .decoding import DecodingOptions, DecodingResult from .tokenizer import LANGUAGES, TO_LANGUAGE_CODE, get_tokenizer -from .utils import exact_div, format_timestamp, optional_int, optional_float, str2bool, write_vtt, write_srt +from .utils import exact_div, format_timestamp, optional_int, optional_float, str2bool, write_txt, write_vtt, write_srt if TYPE_CHECKING: from .model import Whisper @@ -295,7 +295,7 @@ def cli(): # save TXT with open(os.path.join(output_dir, audio_basename + ".txt"), "w", encoding="utf-8") as txt: - print(result["text"], file=txt) + write_txt(result["segments"], file=txt) # save VTT with open(os.path.join(output_dir, audio_basename + ".vtt"), "w", encoding="utf-8") as vtt: diff --git a/whisper/utils.py b/whisper/utils.py index c8f9c0dbe..b63ade7f7 100644 --- a/whisper/utils.py +++ b/whisper/utils.py @@ -44,6 +44,11 @@ def format_timestamp(seconds: float, always_include_hours: bool = False): return f"{hours_marker}{minutes:02d}:{seconds:02d}.{milliseconds:03d}" +def write_txt(transcript: Iterator[dict], file: TextIO): + for segment in transcript: + print(segment['text'].strip(), file=file, flush=True) + + def write_vtt(transcript: Iterator[dict], file: TextIO): print("WEBVTT\n", file=file) for segment in transcript: