Skip to content

Commit

Permalink
Write each sentence as a separate line for the txt output (ggerganov#101
Browse files Browse the repository at this point in the history
)

* 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 <example@example.com>
Co-authored-by: Jong Wook Kim <jongwook@nyu.edu>
  • Loading branch information
3 people authored Sep 26, 2022
1 parent 520796a commit fc0f409
Show file tree
Hide file tree
Showing 2 changed files with 7 additions and 2 deletions.
4 changes: 2 additions & 2 deletions whisper/transcribe.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand Down
5 changes: 5 additions & 0 deletions whisper/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down

0 comments on commit fc0f409

Please sign in to comment.