Merge branch 'main' into srt

README.md

@@ -27,6 +27,9 @@ It also requires the command-line tool [`ffmpeg`](https://ffmpeg.org/) to be ins
 # on Ubuntu or Debian
 sudo apt update && sudo apt install ffmpeg
 
+# on Arch Linux
+sudo pacman -S ffmpeg
+
 # on MacOS using Homebrew (https://brew.sh/)
 brew install ffmpeg
 

whisper/transcribe.py

@@ -5,6 +5,7 @@
 
 import numpy as np
 import torch
+import tqdm
 
 from .audio import SAMPLE_RATE, N_FRAMES, HOP_LENGTH, pad_or_trim, log_mel_spectrogram
 from .decoding import DecodingOptions, DecodingResult
@@ -24,6 +25,7 @@ def transcribe(
  compression_ratio_threshold: Optional[float] = 2.4,
  logprob_threshold: Optional[float] = -1.0,
  no_speech_threshold: Optional[float] = 0.6,
+ condition_on_previous_text: bool = True,
  **decode_options,
 ):
  """
@@ -54,6 +56,11 @@ def transcribe(
  If the no_speech probability is higher than this value AND the average log probability
  over sampled tokens is below `logprob_threshold`, consider the segment as silent
 
+ condition_on_previous_text: bool
+ if True, the previous output of the model is provided as a prompt for the next window;
+ disabling may make the text inconsistent across windows, but the model becomes less prone to
+ getting stuck in a failure loop, such as repetition looping or timestamps going out of sync.
+
  decode_options: dict
  Keyword arguments to construct `DecodingOptions` instances
 
@@ -81,7 +88,7 @@ def transcribe(
  segment = pad_or_trim(mel, N_FRAMES).to(model.device).to(dtype)
  _, probs = model.detect_language(segment)
  decode_options["language"] = max(probs, key=probs.get)
- print(f"Detected language: {LANGUAGES[decode_options['language']]}")
+ print(f"Detected language: {LANGUAGES[decode_options['language']].title()}")
 
  mel = mel.unsqueeze(0)
  language = decode_options["language"]
@@ -154,72 +161,81 @@ def add_segment(
  if verbose:
  print(f"[{format_timestamp(start)} --> {format_timestamp(end)}] {text}")
 
- while seek < mel.shape[-1]:
- timestamp_offset = float(seek * HOP_LENGTH / SAMPLE_RATE)
- segment = pad_or_trim(mel[:, :, seek:], N_FRAMES).to(model.device).to(dtype)
- segment_duration = segment.shape[-1] * HOP_LENGTH / SAMPLE_RATE
-
- decode_options["prompt"] = all_tokens[prompt_reset_since:]
- result = decode_with_fallback(segment)[0]
- tokens = torch.tensor(result.tokens)
-
- if no_speech_threshold is not None:
- # no voice activity check
- should_skip = result.no_speech_prob > no_speech_threshold
- if logprob_threshold is not None and result.avg_logprob > logprob_threshold:
- # don't skip if the logprob is high enough, despite the no_speech_prob
- should_skip = False
-
- if should_skip:
- seek += segment.shape[-1] # fast-forward to the next segment boundary
- continue
-
- timestamp_tokens: torch.Tensor = tokens.ge(tokenizer.timestamp_begin)
- consecutive = torch.where(timestamp_tokens[:-1] & timestamp_tokens[1:])[0].add_(1)
- if len(consecutive) > 0: # if the output contains two consecutive timestamp tokens
- last_slice = 0
- for current_slice in consecutive:
- sliced_tokens = tokens[last_slice:current_slice]
- start_timestamp_position = (
- sliced_tokens[0].item() - tokenizer.timestamp_begin
- )
- end_timestamp_position = (
- sliced_tokens[-1].item() - tokenizer.timestamp_begin
+ # show the progress bar when verbose is False (otherwise the transcribed text will be printed)
+ num_frames = mel.shape[-1]
+ previous_seek_value = seek
+
+ with tqdm.tqdm(total=num_frames, unit='frames', disable=verbose) as pbar:
+ while seek < num_frames:
+ timestamp_offset = float(seek * HOP_LENGTH / SAMPLE_RATE)
+ segment = pad_or_trim(mel[:, :, seek:], N_FRAMES).to(model.device).to(dtype)
+ segment_duration = segment.shape[-1] * HOP_LENGTH / SAMPLE_RATE
+
+ decode_options["prompt"] = all_tokens[prompt_reset_since:]
+ result = decode_with_fallback(segment)[0]
+ tokens = torch.tensor(result.tokens)
+
+ if no_speech_threshold is not None:
+ # no voice activity check
+ should_skip = result.no_speech_prob > no_speech_threshold
+ if logprob_threshold is not None and result.avg_logprob > logprob_threshold:
+ # don't skip if the logprob is high enough, despite the no_speech_prob
+ should_skip = False
+
+ if should_skip:
+ seek += segment.shape[-1] # fast-forward to the next segment boundary
+ continue
+
+ timestamp_tokens: torch.Tensor = tokens.ge(tokenizer.timestamp_begin)
+ consecutive = torch.where(timestamp_tokens[:-1] & timestamp_tokens[1:])[0].add_(1)
+ if len(consecutive) > 0: # if the output contains two consecutive timestamp tokens
+ last_slice = 0
+ for current_slice in consecutive:
+ sliced_tokens = tokens[last_slice:current_slice]
+ start_timestamp_position = (
+ sliced_tokens[0].item() - tokenizer.timestamp_begin
+ )
+ end_timestamp_position = (
+ sliced_tokens[-1].item() - tokenizer.timestamp_begin
+ )
+ add_segment(
+ start=timestamp_offset + start_timestamp_position * time_precision,
+ end=timestamp_offset + end_timestamp_position * time_precision,
+ text_tokens=sliced_tokens[1:-1],
+ result=result,
+ )
+ last_slice = current_slice
+ last_timestamp_position = (
+ tokens[last_slice - 1].item() - tokenizer.timestamp_begin
  )
+ seek += last_timestamp_position * input_stride
+ all_tokens.extend(tokens[: last_slice + 1].tolist())
+ else:
+ duration = segment_duration
+ timestamps = tokens[timestamp_tokens.nonzero().flatten()]
+ if len(timestamps) > 0:
+ # no consecutive timestamps but it has a timestamp; use the last one.
+ # single timestamp at the end means no speech after the last timestamp.
+ last_timestamp_position = timestamps[-1].item() - tokenizer.timestamp_begin
+ duration = last_timestamp_position * time_precision
+
  add_segment(
- start=timestamp_offset + start_timestamp_position * time_precision,
- end=timestamp_offset + end_timestamp_position * time_precision,
- text_tokens=sliced_tokens[1:-1],
+ start=timestamp_offset,
+ end=timestamp_offset + duration,
+ text_tokens=tokens,
  result=result,
  )
- last_slice = current_slice
- last_timestamp_position = (
- tokens[last_slice - 1].item() - tokenizer.timestamp_begin
- )
- seek += last_timestamp_position * input_stride
- all_tokens.extend(tokens[: last_slice + 1].tolist())
- else:
- duration = segment_duration
- timestamps = tokens[timestamp_tokens.nonzero().flatten()]
- if len(timestamps) > 0:
- # no consecutive timestamps but it has a timestamp; use the last one.
- # single timestamp at the end means no speech after the last timestamp.
- last_timestamp_position = timestamps[-1].item() - tokenizer.timestamp_begin
- duration = last_timestamp_position * time_precision
-
- add_segment(
- start=timestamp_offset,
- end=timestamp_offset + duration,
- text_tokens=tokens,
- result=result,
- )
-
- seek += segment.shape[-1]
- all_tokens.extend(tokens.tolist())
-
- if result.temperature > 0.5:
- # do not feed the prompt tokens if a high temperature was used
- prompt_reset_since = len(all_tokens)
+
+ seek += segment.shape[-1]
+ all_tokens.extend(tokens.tolist())
+
+ if not condition_on_previous_text or result.temperature > 0.5:
+ # do not feed the prompt tokens if a high temperature was used
+ prompt_reset_since = len(all_tokens)
+
+ # update progress bar
+ pbar.update(min(num_frames, seek) - previous_seek_value)
+ previous_seek_value = seek
 
  return dict(text=tokenizer.decode(all_tokens), segments=all_segments, language=language)
 
@@ -232,7 +248,7 @@ def cli():
  parser.add_argument("--model", default="small", choices=available_models(), help="name of the Whisper model to use")
  parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu", help="device to use for PyTorch inference")
  parser.add_argument("--output_dir", "-o", type=str, default=".", help="directory to save the outputs")
- parser.add_argument("--verbose", type=str2bool, default=True, help="Whether to print out the progress and debug messages")
+ parser.add_argument("--verbose", type=str2bool, default=True, help="whether to print out the progress and debug messages")
 
  parser.add_argument("--task", type=str, default="transcribe", choices=["transcribe", "translate"], help="whether to perform X->X speech recognition ('transcribe') or X->English translation ('translate')")
  parser.add_argument("--language", type=str, default=None, choices=sorted(LANGUAGES.keys()) + sorted([k.title() for k in TO_LANGUAGE_CODE.keys()]), help="language spoken in the audio, specify None to perform language detection")
@@ -244,6 +260,7 @@ def cli():
  parser.add_argument("--length_penalty", type=float, default=None, help="optional token length penalty coefficient (alpha) as in https://arxiv.org/abs/1609.08144, uses simple lengt normalization by default")
 
  parser.add_argument("--suppress_tokens", type=str, default="-1", help="comma-separated list of token ids to suppress during sampling; '-1' will suppress most special characters except common punctuations")
+ parser.add_argument("--condition_on_previous_text", type=str2bool, default=True, help="if True, provide the previous output of the model as a prompt for the next window; disabling may make the text inconsistent across windows, but the model becomes less prone to getting stuck in a failure loop")
  parser.add_argument("--fp16", type=str2bool, default=True, help="whether to perform inference in fp16; True by default")
 
  parser.add_argument("--temperature_increment_on_fallback", type=optional_float, default=0.2, help="temperature to increase when falling back when the decoding fails to meet either of the thresholds below")
@@ -257,7 +274,7 @@ def cli():
  device: str = args.pop("device")
  os.makedirs(output_dir, exist_ok=True)
 
- if model_name.endswith(".en") and args["language"] != "en":
+ if model_name.endswith(".en") and args["language"] not in {"en", "English"}:
  warnings.warn(f"{model_name} is an English-only model but receipted '{args['language']}'; using English instead.")
  args["language"] = "en"