Spaces:

RaeesRSB
/

word-sub

Runtime error

App Files Files Community

word-sub / whisper_timestamped /transcribe.py

RaeesRSB

Upload 185 files

2cba4ca 6 months ago

raw history blame contribute delete

No virus

123 kB

	#!/usr/bin/env python3

	__author__ = "Jérôme Louradour"
	__credits__ = ["Jérôme Louradour"]
	__license__ = "GPLv3"
	__version__ = "1.14.2"

	# Set some environment variables
	import os
	os.environ['TF_CPP_MIN_LOG_LEVEL'] = '1' # Remove warning "This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN)..."
	os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID' # GPU in the right order

	# Whisper and Torch
	import whisper
	import torch
	import torch.nn.functional as F

	from importlib.util import find_spec
	if find_spec("intel_extension_for_pytorch") is not None:
	try:
	import intel_extension_for_pytorch
	except ImportError:
	pass

	# For alignment
	import numpy as np
	import dtw
	# from scipy.signal import medfilt as median_filter
	from scipy.ndimage import median_filter # faster owing to https://github.com/openai/whisper/commit/f0083e7eb20d032390e42f6f6039947fa8669c93
	from scipy.signal import find_peaks

	# Additional
	import string
	import csv
	import sys
	import gzip, base64
	import copy
	import re
	import shutil

	# Constant variables
	from whisper.utils import format_timestamp
	from whisper.audio import N_FRAMES, HOP_LENGTH, SAMPLE_RATE # 3000, 160, 16000
	AUDIO_SAMPLES_PER_TOKEN = HOP_LENGTH * 2 # 320
	AUDIO_TIME_PER_TOKEN = AUDIO_SAMPLES_PER_TOKEN / SAMPLE_RATE # 0.02 (sec)
	SEGMENT_DURATION = N_FRAMES * HOP_LENGTH / SAMPLE_RATE # 30.0 (sec)

	# Logs
	import logging
	logger = logging.getLogger("whisper_timestamped")

	USE_EFFICIENT_BY_DEFAULT = True
	TRUST_WHISPER_TIMESTAMP_BY_DEFAULT = True
	DISFLUENCY_MARK = "[*]"

	try:
	whisper_version = whisper.__version__
	except NameError:
	whisper_version = ""
	WHIPSER_GE_20230306 = whisper_version >= "20230306"
	WHIPSER_GE_20230308 = whisper_version >= "20230308"

	def transcribe_timestamped(
	# Main Whisper options
	model,
	audio,
	language=None,
	task="transcribe",

	# Additional options for word alignment
	remove_punctuation_from_words=False,
	compute_word_confidence=True,
	include_punctuation_in_confidence=False,
	refine_whisper_precision=0.5,
	min_word_duration=0.02, # Was 0.04 before 1.11
	plot_word_alignment=False,
	word_alignement_most_top_layers=None, # Was 6 before 1.9
	remove_empty_words=False,

	# Reproducibility
	seed=1234,

	vad=False,
	detect_disfluencies=False,
	trust_whisper_timestamps=TRUST_WHISPER_TIMESTAMP_BY_DEFAULT,
	naive_approach=False,

	# Other Whisper options
	temperature=0.0 if USE_EFFICIENT_BY_DEFAULT else (0.0, 0.2, 0.4, 0.6, 0.8, 1.0),
	best_of=None,
	beam_size=None,
	patience=None,
	length_penalty=None,
	compression_ratio_threshold=2.4,
	logprob_threshold=-1.0,
	no_speech_threshold=0.6,
	fp16=None,
	condition_on_previous_text=True,
	initial_prompt=None,
	suppress_tokens="-1",
	sample_len=None,
	verbose=False,
	):
	"""
	Transcribe an audio file using Whisper

	Parameters
	----------
	model: Whisper
	The Whisper model instance.

	audio: Union[str, np.ndarray, torch.Tensor]
	The path to the audio file to open, or the audio waveform in 16kHz.

	language: str
	The language to use for the transcription. If None, the language is detected automatically.

	task: str
	The task to perform: either "transcribe" or "translate".

	remove_punctuation_from_words: bool
	If False, words will be glued with the next punctuation mark (if any).
	If True, there will be no punctuation mark in the `words[:]["text"]` list.
	It only affects these strings; This has no influence on the computation of the word confidence, whatever the value of `include_punctuation_in_confidence` is.

	include_punctuation_in_confidence: bool
	Whether to include proba of punctuation in the computation of the (previous) word confidence.

	compute_word_confidence: bool
	Whether to compute word confidence.
	If True, a finer confidence for each segment will be computed as well.

	vad: bool or str in ["silero", "silero:3.1", "auditok"]
	Whether to perform voice activity detection (VAD) on the audio file, to remove silent parts before transcribing with Whisper model.
	This should decrease hallucinations from the Whisper model.
	When set to True, the default VAD algorithm is used (silero).
	When set to a string, the corresponding VAD algorithm is used (silero, silero:3.1 or auditok).
	Note that the library for the corresponding VAD algorithm must be installed.

	detect_disfluencies: bool
	Whether to detect disfluencies (i.e. hesitations, filler words, repetitions, corrections, etc.) that Whisper model might have omitted in the transcription.
	This should make the word timestamp prediction more accurate.
	And probable disfluencies will be marked as special words "[*]".

	trust_whisper_timestamps: bool
	Whether to rely on Whisper's timestamps to get approximative first estimate of segment positions (up to refine_whisper_precision).

	refine_whisper_precision: float
	How much can we refine Whisper segment positions, in seconds. Must be a multiple of 0.02.

	min_word_duration: float
	Minimum duration of a word, in seconds. If a word is shorter than this, timestamps will be adjusted.

	plot_word_alignment: bool
	Whether to plot the word alignment for each segment. matplotlib must be installed to use this option.

	remove_empty_words: bool
	Whether to remove words with no duration occuring at the end of segments (probable Whisper hallucinations).

	seed: int
	Random seed to use for temperature sampling, for the sake of reproducibility.
	Choose None for unpredictable randomness.

	naive_approach: bool
	Force the naive approach that consists in decoding twice the audio file, once to get the transcritpion and once with the decoded tokens to get the alignment.
	Note that this approach is used anyway when beam_size is not None and/or when the temperature is a list with more than one element.

	temperature: float
	Temperature for sampling.

	compression_ratio_threshold: float
	If the gzip compression ratio is above this value, treat as failed.

	logprob_threshold: float
	If the average log probability over sampled tokens is below this value, treat as failed.

	no_speech_threshold: float
	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.

	initial_prompt: str
	Optional text to provide as a prompt for the first window.

	suppress_tokens: str
	Comma-separated list of token ids to suppress during sampling;
	'-1' will suppress most special characters except common punctuations.

	verbose: bool
	Whether to display the text being decoded to the console. If True, displays all the details,
	If False, displays minimal details. If None, does not display anything

	Returns
	-------
	A dictionary containing the resulting text ("text") and segment-level details ("segments"), and
	the spoken language ("language"), which is detected when `decode_options["language"]` is None.
	"""

	if seed is not None:
	torch.manual_seed(seed)
	torch.cuda.manual_seed_all(seed)

	# Check input options
	assert refine_whisper_precision >= 0 and refine_whisper_precision / AUDIO_TIME_PER_TOKEN == round(refine_whisper_precision / AUDIO_TIME_PER_TOKEN), f"refine_whisper_precision must be a positive multiple of {AUDIO_TIME_PER_TOKEN}"
	refine_whisper_precision_nframes = round(refine_whisper_precision / AUDIO_TIME_PER_TOKEN)
	assert min_word_duration >= 0, f"min_word_duration must be a positive number"
	assert word_alignement_most_top_layers is None or word_alignement_most_top_layers > 0, f"word_alignement_most_top_layers must be a strictly positive number"

	if isinstance(temperature, (list, tuple)) and len(temperature) == 1:
	temperature = temperature[0]
	if isinstance(temperature, (list, tuple)):
	# temperature fallback
	naive_approach = True
	elif temperature > 0 and best_of is not None and best_of > 1:
	naive_approach = True
	if beam_size is not None:
	# beam-search
	naive_approach = True

	# Input options
	vad = check_vad_method(vad)
	if isinstance(model, str):
	model = load_model(model)
	if fp16 is None:
	fp16 = model.device != torch.device("cpu")

	# Safety check
	input_stride = N_FRAMES // model.dims.n_audio_ctx
	time_precision = input_stride * HOP_LENGTH / SAMPLE_RATE
	assert time_precision == AUDIO_TIME_PER_TOKEN

	alignment_heads = get_alignment_heads(model) if word_alignement_most_top_layers is None else None
	if alignment_heads is None and word_alignement_most_top_layers is None:
	word_alignement_most_top_layers = 6

	alignment_options = dict(
	remove_punctuation_from_words=remove_punctuation_from_words,
	compute_word_confidence=compute_word_confidence,
	include_punctuation_in_confidence=include_punctuation_in_confidence,
	detect_disfluencies=detect_disfluencies,
	refine_whisper_precision_nframes=refine_whisper_precision_nframes,
	plot_word_alignment=plot_word_alignment,
	word_alignement_most_top_layers=word_alignement_most_top_layers,
	alignment_heads=alignment_heads,
	)
	whisper_options = dict(
	language=language,
	task=task,
	fp16=fp16,
	temperature=temperature,
	best_of=best_of,
	beam_size=beam_size,
	patience=patience,
	length_penalty=length_penalty,
	condition_on_previous_text=condition_on_previous_text,
	initial_prompt=initial_prompt,
	suppress_tokens=suppress_tokens,
	sample_len=sample_len,
	verbose=verbose if (not vad or verbose is not True) else False,
	)
	other_options = dict(
	no_speech_threshold=no_speech_threshold,
	logprob_threshold=logprob_threshold,
	compression_ratio_threshold=compression_ratio_threshold,
	)

	if vad:
	audio = get_audio_tensor(audio)
	audio, convert_timestamps = remove_non_speech(audio, method=vad, plot=plot_word_alignment)

	global num_alignment_for_plot
	num_alignment_for_plot = 0

	if naive_approach:
	(transcription, words) = _transcribe_timestamped_naive(model, audio,
	min_word_duration=0.0, # Was 0.04 before 1.11
	trust_whisper_timestamps=trust_whisper_timestamps,
	alignment_options, whisper_options, **other_options)
	else:
	(transcription, words) = _transcribe_timestamped_efficient(model, audio,
	trust_whisper_timestamps=trust_whisper_timestamps,
	alignment_options, whisper_options, **other_options)
	if remove_empty_words:
	# Remove words with empty duration happening at the end of segments, to remove some hallucinations
	transcription, words = remove_last_null_duration_words(transcription, words, recompute_text=True)

	# Refine word positions
	ensure_increasing_positions(words, min_duration=min_word_duration if trust_whisper_timestamps else 0)

	# Combine words and segments
	whisper_segments = transcription["segments"]
	for word in words:
	if verbose and not naive_approach and not vad:
	print_timestamped(word)
	word.pop("tokens")
	word.pop("tokens_indices")
	if "avg_logprob_reliable" in word:
	word.pop("avg_logprob_reliable")
	idx_segment = word.pop("idx_segment")
	assert idx_segment < len(whisper_segments), f"Fatal error: Got unexpected segment index {idx_segment} >= {len(whisper_segments)}"
	segment = whisper_segments[idx_segment]
	if "words" in segment:
	segment["words"].append(word)
	else:
	segment["words"] = [word]
	if refine_whisper_precision:
	segment["start"] = word["start"]
	if refine_whisper_precision:
	segment["end"] = word["end"]

	if vad:
	# Recompute timestamps to match the original audio
	for segment in whisper_segments:
	for word in segment.get("words", []):
	word["start"], word["end"] = convert_timestamps(word["start"], word["end"])
	if verbose:
	print_timestamped(word)
	if refine_whisper_precision and len(segment.get("words", [])):
	segment["start"] = segment["words"][0]["start"]
	segment["end"] = segment["words"][-1]["end"]
	else:
	segment["start"], segment["end"] = convert_timestamps(segment["start"], segment["end"])

	return transcription

	def _transcribe_timestamped_efficient(
	model,
	audio,
	remove_punctuation_from_words,
	compute_word_confidence,
	include_punctuation_in_confidence,
	refine_whisper_precision_nframes,
	alignment_heads,
	plot_word_alignment,
	word_alignement_most_top_layers,
	detect_disfluencies,
	trust_whisper_timestamps,
	use_timestamps_for_alignment = True,
	# Whisper specific options
	**whisper_options,
	):

	# Get options
	sample_len = whisper_options["sample_len"]
	temperature = whisper_options["temperature"]
	no_speech_threshold = whisper_options["no_speech_threshold"]
	logprob_threshold = whisper_options["logprob_threshold"]
	verbose = whisper_options["verbose"]
	# Note: "on-the-fly" verbose is not implementable in the current state (we don't know the absolute position of the current chunk). See issue #18
	verbose_bugged = False
	whisper_options["verbose"] = None if whisper_options["verbose"] is True else whisper_options["verbose"] # We will print intermediate results ourselves

	logit_filters = get_logit_filters(model, whisper_options)
	language = whisper_options["language"]
	tokenizer = get_tokenizer(model, task=whisper_options["task"], language=language)

	max_sample_len = sample_len or model.dims.n_text_ctx // 2
	n_ctx = model.dims.n_text_ctx

	debug = logger.getEffectiveLevel() >= logging.DEBUG

	word_alignement_most_top_layers = float("inf") if word_alignement_most_top_layers is None else word_alignement_most_top_layers

	# The main outcome
	timestamped_word_segments = [] # list of timestamped word segments that have been collected so far
	# Main variables to be accumulated
	segment_tokens = [[]] # list of lists of token indices that have been collected so far (one list per segment)
	segment_attweights = [[] for _ in range(min(word_alignement_most_top_layers, len(model.decoder.blocks)))]
	# attention weights on the last segments
	segment_avglogprobs = [] # average log probability for each segment (actually of the corresponding chunk, as computed by whisper)
	segment_logprobs = [] # token log probabilities for each segment
	# Variables related to options that can skip some segments
	sot_index = None # index of the SOT token in the current set of processed tokens
	no_speech_prob = None # no speech probability for the current 30 sec chunk
	chunk_logprobs = [] # log probabilities for the current 30 sec chunk
	chunk_tokens = [] # tokens for the current 30 sec chunk (list of Torch tensors)
	chunk_tokens_nosot = [] # tokens for the current 30 sec chunk, without the SOT tokens (list of indices)
	last_chunk_token = None # last token of the current chunk, that may be needed for corner cases
	last_token_fallback = None # last token to use as a fallback if the model gets stuck
	has_started = False # whether we have started decoding
	mfcc = None # MFCC features for the current 30 sec chunk
	new_mfcc = None #
	num_inference_steps = 0 # number of inference steps performed so far (for debugging only)
	language_probs = None # language detection probabilities

	def is_sot(curr_tokens):
	return curr_tokens is None or len(curr_tokens) > 1 or curr_tokens[0] == tokenizer.sot

	def has_reached_decoding_limit():
	n = len(chunk_tokens_nosot) + 1
	m = n + (len(chunk_tokens[0]) if len(chunk_tokens) > 0 else 0)
	return n + 1 >= max_sample_len or m > n_ctx

	def reset(add_segment, keep_last_token=True):
	""" Reset the list of tokens for the current speech segment, and corresponding cross-attention weights """
	nonlocal segment_tokens, segment_attweights
	if add_segment:
	if keep_last_token:
	segment_tokens.append([segment_tokens[-1][-1]])
	segment_attweights = [w[-1:] for w in segment_attweights]
	else:
	segment_tokens.append([])
	segment_attweights = [[] for w in segment_attweights]
	segment_tokens[-2].pop(0)
	elif len(segment_tokens[-1]) > 0:
	if debug:
	logger.debug(f"Reset last segment: {tokenizer.decode_with_timestamps(segment_tokens[-1])}")
	segment_tokens[-1] = []
	segment_attweights = [[] for w in segment_attweights]

	saw_consecutive_timestamps = False
	def must_flush_segment(curr_tokens):
	""" Return whether or not the previously collected tokens must be used to add a new speech segment """
	nonlocal segment_tokens, saw_consecutive_timestamps, chunk_tokens_nosot

	if not is_sot(curr_tokens):
	is_timestamp = curr_tokens[0] >= tokenizer.timestamp_begin
	is_previous_timestamp = segment_tokens[-1][-1] >= tokenizer.timestamp_begin if len(segment_tokens[-1]) > 0 else False
	consecutive_timestamps = is_timestamp and is_previous_timestamp
	if consecutive_timestamps:
	saw_consecutive_timestamps = True
	return consecutive_timestamps
	else: # Several tokens as a prompt or must flush last segments

	must_flush = len(segment_tokens[-1]) > 1 and not saw_consecutive_timestamps
	if not must_flush and WHIPSER_GE_20230306: # If the last token is a timestamp, the last segment is used
	if last_chunk_token is None:
	must_flush = (len(segment_tokens[-1]) > 2 and segment_tokens[-1][-1] >= tokenizer.timestamp_begin)
	else:
	must_flush = (last_chunk_token >= tokenizer.timestamp_begin)
	if not must_flush and trust_whisper_timestamps:
	# Discard the end of the last transcription
	reset(False)
	saw_consecutive_timestamps = False
	return must_flush

	index_begin_30sec_chunck = 0
	def get_index_begin_30sec_chunck(curr_tokens):
	nonlocal index_begin_30sec_chunck, has_started

	if is_sot(curr_tokens) and has_started:
	if trust_whisper_timestamps:
	res = index_begin_30sec_chunck
	index_begin_30sec_chunck = len(segment_tokens)-1
	else:
	res = len(segment_tokens)-1
	return res

	def align_last_segment(curr_tokens=None):
	nonlocal segment_tokens, segment_attweights, timestamped_word_segments, has_started, no_speech_prob, chunk_tokens, chunk_tokens_nosot, chunk_logprobs, mfcc, new_mfcc, logit_filters, index_begin_30sec_chunck, last_token_fallback, num_inference_steps

	if debug and trust_whisper_timestamps:
	logger.debug(f"Add segment {len(timestamped_word_segments)+1} at step {num_inference_steps}:\n\t{tokenizer.decode_with_timestamps(segment_tokens[-1])}")

	tokens = segment_tokens[-1][1:]

	# When the decoding hit the max limit (number of tokens) -- usually when the language model gets stuck --
	# then we have to recover the last token from what is send to the decoder
	unfinished_decoding = has_reached_decoding_limit()
	last_is_not_timestamp = len(tokens) and tokens[-1] < tokenizer.timestamp_begin
	last_token_reliable = True

	if unfinished_decoding:
	logger.debug(f"WARNING: decoding hit the max limit for segment {segment_tokens[-1]} (It usually happens when the language model gets stuck)")
	# The last token chosen is in the prompt for the new chunk
	if curr_tokens is not None and curr_tokens[0] == tokenizer.sot_prev:
	index_sot = (curr_tokens == tokenizer.sot).nonzero(as_tuple=True)
	assert len(index_sot) == 1
	index_sot = index_sot[0].item()
	assert index_sot > 0
	last_token_fallback = curr_tokens[index_sot-1].item()
	logger.debug(f" Guessed last token from the prompt for the new chunk: {last_token_fallback}")
	# Fallback for the last segment, or without prompt: Assume greedy decoding
	else:
	last_token_fallback = torch.argmax(chunk_logprobs[-1]).item() if last_chunk_token is None else last_chunk_token
	last_token_reliable = (temperature == 0)
	logger.debug(f" Guess last token using probas (assuming greedy decoding): {last_token_fallback}")
	if debug:
	logger.debug(f"WARNING: also add last token: {tokenizer.decode_with_timestamps([last_token_fallback])}")

	tokens.append(last_token_fallback)
	segment_tokens[-1].append(last_token_fallback)
	attention_weights = [torch.cat(w, dim=-2) for w in segment_attweights]
	last_logprobs = chunk_logprobs[-1]
	elif last_is_not_timestamp: # <eot> was emitted early, without a timestamp before
	logger.debug(f"WARNING: end timestamp not produced. Adding <\|endoftext\|>")
	tokens.append(tokenizer.eot)
	segment_tokens[-1].append(tokenizer.eot)
	attention_weights = [torch.cat(w, dim=-2) for w in segment_attweights]
	last_logprobs = chunk_logprobs[-1]
	else:
	attention_weights = [torch.cat(w[:-1], dim=-2) for w in segment_attweights]
	last_logprobs = chunk_logprobs[-2]

	# Check prediction of last token
	end_token = tokens[-1]
	if end_token >= tokenizer.timestamp_begin:
	start_token = tokens[0]
	assert start_token >= tokenizer.timestamp_begin
	# If Whisper prediction of the end is obviously wrong, we predict it again (constrained)
	if end_token <= start_token:
	new_end_token = last_logprobs[start_token+1:].argmax() + start_token + 1
	tokens[-1] = new_end_token.item()
	if debug:
	logger.debug(f"Re-estimated end token {tokenizer.decode_with_timestamps([new_end_token])} (was {tokenizer.decode_with_timestamps([end_token])}) to be after start token {tokenizer.decode_with_timestamps([start_token])}")

	if len(tokens) <= 1:
	# Corner case: nothing in between timestamps
	ws = []
	else:
	ws = perform_word_alignment(
	tokens,
	attention_weights,
	tokenizer,
	use_space=should_use_space(language),
	alignment_heads=alignment_heads,
	remove_punctuation_from_words=remove_punctuation_from_words,
	refine_whisper_precision_nframes=refine_whisper_precision_nframes,
	detect_disfluencies=detect_disfluencies,
	unfinished_decoding=unfinished_decoding,
	mfcc=mfcc,
	plot=plot_word_alignment,
	debug=debug,
	)

	add_segment = len(ws) > 0
	if add_segment:
	timestamped_word_segments.append(ws)
	else:
	logger.debug(f"Not added!")
	reset(add_segment, not is_sot(curr_tokens))

	return add_segment, unfinished_decoding, last_token_reliable

	def may_flush_segment(curr_tokens = None):
	""" Add a speech segment with the new tokens if necessary.
	May also remove the last collected segments if filtered out by Whisper (no_speech_prob <= no_speech_threshold)
	"""
	nonlocal segment_tokens, segment_attweights, timestamped_word_segments, segment_logprobs, has_started, no_speech_prob, chunk_tokens, chunk_tokens_nosot, chunk_logprobs, mfcc, new_mfcc, logit_filters, index_begin_30sec_chunck, last_token_fallback, num_inference_steps, last_chunk_token

	# Check if a new segment should be added
	unfinished_decoding = False
	last_token_reliable = True

	if must_flush_segment(curr_tokens) and trust_whisper_timestamps:
	_, unfinished_decoding, last_token_reliable = align_last_segment(curr_tokens)

	i_start = get_index_begin_30sec_chunck(curr_tokens)

	# All segments from previous 30sec chunck have been collected
	if i_start is not None:

	if not trust_whisper_timestamps:

	tokens = torch.Tensor(segment_tokens[-1]).int()
	idx_task = torch.where(tokens==tokenizer.sot_sequence[-1])[0][0].item() # index of <\|transcribe\|>

	is_special = tokens.ge(tokenizer.eot)
	# Remove prompt
	is_special[:idx_task] = True
	# Keep begin timestamp
	is_special[idx_task:idx_task+2] = False

	is_timestamp = tokens.ge(tokenizer.timestamp_begin)
	consecutive = torch.where(is_timestamp[1:] & is_timestamp[:-1])[0]
	if (WHIPSER_GE_20230306 or has_reached_decoding_limit()) and (
	(is_timestamp[-1] and not is_timestamp[-2]) if last_chunk_token is None else
	last_chunk_token >= tokenizer.timestamp_begin and not is_timestamp[-2]
	):
	consecutive = torch.cat([consecutive, torch.Tensor([len(tokens)-1]).int()])
	last_is_timestamp = True
	if len(consecutive):
	# Remove last tokens
	is_special[consecutive[-1]+1:] = True
	# Keep end timestamp
	is_special[consecutive[-1]] = False
	elif is_timestamp[-1]:
	# Keep end timestamp
	is_special[-1] = False
	else:
	last_is_timestamp = False

	if use_timestamps_for_alignment and len(consecutive):
	# Keep all timestamps
	is_special[idx_task+2:consecutive[-1]] = False

	# Do remove what has to be removed
	is_next_achar = ~torch.cat([is_special[1:], torch.Tensor([False]).bool()])
	for i, weights in enumerate(segment_attweights):
	assert len(weights) == len(tokens), f"{len(weights)} attention weights != {len(tokens)}"
	# We must remove attention weights used to predict timestamp tokens
	segment_attweights[i] = [w for s, w in zip(is_next_achar, weights) if s]
	tokens_filtered = tokens[~is_special]
	assert len(segment_attweights[0]) == len(tokens_filtered), f"{len(segment_attweights[0])} attention weights != {len(tokens_filtered)} "

	# Replace first and last timestamp
	orig_start, orig_end = tokens_filtered[1].item(), tokens_filtered[-1].item()
	tokens_filtered[1] = tokenizer.timestamp_begin # <\|0.00\|>
	if last_is_timestamp:
	tokens_filtered[-1] = tokenizer.timestamp_begin + N_FRAMES // 2 # <\|30.00\|>
	segment_tokens[-1] = tokens_filtered.tolist()

	# Do alignement
	added, unfinished_decoding, last_token_reliable = align_last_segment()

	# Re-split into segments (if necessary)
	if added:
	if len(consecutive) > 1:
	segments_timestamped_concat = timestamped_word_segments[-1]
	new_segments_timestamped = []
	new_segment_tokens = []
	start = idx_task+1
	i_word = 0
	for i, end in enumerate(consecutive):
	end = end.item()
	new_segment_tokens.append(tokens[start:end+1].tolist())
	if debug:
	logger.debug(f"Add segment {len(timestamped_word_segments)+i}:\n\t{tokenizer.decode_with_timestamps(new_segment_tokens[-1])}")
	total_length = end - start - 1
	start = end+1
	length = 0
	new_segments_timestamped.append([])
	while length < total_length:
	if not use_timestamps_for_alignment and i_word == len(segments_timestamped_concat):
	# This can happen in the case of "..."
	assert total_length == 1 and i == len(consecutive)-1, "Unexpected situation!"
	break
	assert i_word < len(segments_timestamped_concat), f"i_word={i_word} < len(segments_timestamped_concat)={len(segments_timestamped_concat)}"
	word = segments_timestamped_concat[i_word]
	new_segments_timestamped[-1].append(word)
	length += len(word["tokens_indices"])
	i_word += 1
	# This can be non zero, when a punctuation (alone in a segment) is glued to the previous segment
	if use_timestamps_for_alignment:
	assert length == total_length, f"length={length} != total_length={total_length}"
	elif length > total_length:
	delta = length - total_length
	word = new_segments_timestamped[-1][-1]
	word_tokindices = word["tokens_indices"]
	word_tokens = word["tokens"]
	word["tokens_indices"] = word_tokindices[:-delta]
	word["tokens"] = word_tokens[:-delta]
	word["word"] = "".join(word_tokens[:-delta])
	i_word -= 1
	t = segments_timestamped_concat[i_word]["end"]
	segments_timestamped_concat[i_word] = dict(
	text="".join(word_tokens[-delta:]),
	start=t, end=t, # Word without timestamp
	tokens=word_tokens[-delta:],
	tokens_indices=word_tokindices[-delta:],
	)

	assert i_word == len(segments_timestamped_concat)

	segment_tokens = segment_tokens[:-2] + new_segment_tokens + [segment_tokens[-1]]
	timestamped_word_segments = timestamped_word_segments[:-1] + new_segments_timestamped

	else:

	# Recover start and end token
	segment = segment_tokens[-2]
	tokenizer.decode_with_timestamps([orig_start,orig_end])
	segment[0] = orig_start
	if last_is_timestamp:
	segment[-1] = orig_end

	if debug:
	logger.debug(f"Add segment {len(timestamped_word_segments)}:\n\t{tokenizer.decode_with_timestamps(segment)}")

	if unfinished_decoding:
	timestamped_word_segments[-1][-1]["avg_logprob_reliable"] = last_token_reliable

	reset(False)

	mfcc = new_mfcc

	n_segments = len(segment_tokens)-1

	# Get word confidence and/or check if previous segments shoud have been skipped
	should_skip = False
	if compute_word_confidence or no_speech_threshold is not None:

	# no voice activity check
	should_skip = (no_speech_prob > no_speech_threshold) if (no_speech_threshold is not None) else False
	if compute_word_confidence or (should_skip and logprob_threshold is not None):
	n = len(chunk_logprobs)
	if n == len(chunk_tokens_nosot):
	chunk_tokens_nosot = chunk_tokens_nosot[1:]
	if unfinished_decoding:
	assert last_token_fallback is not None
	last_tokens = [last_token_fallback]
	timestamped_word_segments[-1][-1]["avg_logprob_reliable"] = last_token_reliable
	n += 1
	elif has_reached_decoding_limit():
	# there were segments in the 30sec chunck, and then the LM got stuck
	last_tokens = [torch.argmax(chunk_logprobs[-1]).item()]
	timestamped_word_segments[-1][-1]["avg_logprob_reliable"] = (temperature == 0)
	else:
	last_tokens = [tokenizer.eot]
	chunck_indices = chunk_tokens_nosot + last_tokens
	assert len(chunk_logprobs) == len(chunck_indices), f"{len(chunk_logprobs)} != {len(chunck_indices)}"
	logprobs = torch.cat([logprob[i].unsqueeze(0) for (logprob, i) in zip(chunk_logprobs, chunck_indices)])
	assert min([p.isfinite().item() for p in logprobs]), \
	f"Got infinite logprob among ({len(logprobs)}) {[(i, tokenizer.decode_with_timestamps([i]), v.item()) for (i,v) in zip(chunck_indices, logprobs)]}"
	sum_logprob = sum(logprobs)
	avg_logprob = sum_logprob/n
	# don't skip if the logprob is high enough, whatever the no_speech_prob is
	if logprob_threshold is not None and avg_logprob > logprob_threshold:
	should_skip = False

	if should_skip:
	logger.debug(f"Skipping last {n_segments-i_start} segments (no_speech_prob {no_speech_prob} > {no_speech_threshold} and avg_logprob {avg_logprob} < {logprob_threshold})")
	index_begin_30sec_chunck -= n_segments-i_start
	segment_tokens = segment_tokens[:i_start] + [segment_tokens[-1]]
	timestamped_word_segments = timestamped_word_segments[:i_start]
	elif compute_word_confidence:
	avg_logprob = avg_logprob.item()
	i_token_end = -1
	for i in range(i_start, n_segments):
	tokens = segment_tokens[i]
	i_token_start = i_token_end + 1
	i_token_end = i_token_start + len(tokens)
	assert chunck_indices[i_token_start:i_token_end] == tokens, f"Inconsistent token list {tokenizer.decode_with_timestamps(chunck_indices[i_token_start:i_token_end])} != {tokenizer.decode_with_timestamps(tokens)}"
	i_token_start += 1 # skip sos (start time)
	if not unfinished_decoding or i != n_segments-1:
	i_token_end -= 1 # skip eos (end time)
	segment_logprobs.append(logprobs[i_token_start:i_token_end])
	segment_avglogprobs.append(avg_logprob)
	else:
	for i in range(i_start, n_segments):
	segment_logprobs.append(None)
	segment_avglogprobs.append(None)

	else:
	for i in range(i_start, n_segments):
	segment_logprobs.append(None)
	segment_avglogprobs.append(None)

	if verbose_bugged and not should_skip:
	for segment in timestamped_word_segments[i_start:]:
	for word in segment:
	print_timestamped(word)

	# Reset counters
	chunk_tokens = []
	chunk_tokens_nosot = []
	chunk_logprobs = []
	no_speech_prob = None

	def hook_attention_weights(layer, ins, outs, index):
	nonlocal segment_attweights
	# In old version of whisper, output is a single tensor
	assert isinstance(outs, tuple) and len(outs) == 2, "whisper seems to be outdated, please update it (pip install --upgrade --no-deps --force-reinstall git+https://github.com/openai/whisper.git)"
	if not has_started:
	return
	w = outs[-1]
	# Only the last attention weights is useful
	if w.shape[-2] > 1:
	w = w[:, :, -1:, :]
	segment_attweights[index].append(w.cpu())

	def hook_mfcc(layer, ins, outs):
	nonlocal new_mfcc, mfcc
	new_mfcc = ins[0]
	if mfcc is None:
	mfcc = new_mfcc

	def hook_input_tokens(layer, ins, outs):
	nonlocal segment_tokens, sot_index, chunk_tokens, chunk_tokens_nosot, logit_filters, has_started, language, num_inference_steps
	num_inference_steps += 1

	curr_tokens = ins[0]
	assert curr_tokens.shape[0] == 1, "Batch decoding is not supported"
	curr_tokens = curr_tokens.squeeze(0)

	if is_sot(curr_tokens):
	chunk_prompt = curr_tokens.tolist()
	if language is None:
	if len(curr_tokens) > 1:
	language = tokenizer.decode(curr_tokens[-2:-1])
	language = language[2:-2] # remove trailing "<\|" and "\|>"
	whisper_options["language"] = language

	if verbose and not whisper_options["verbose"] and len(curr_tokens) > 1:
	# Reproduce whisper verbose (2/2)
	print(f"Detected language: {whisper.tokenizer.LANGUAGES[language].title()}")
	sys.stdout.flush()

	logit_filters = get_logit_filters(model, whisper_options, prompt = chunk_prompt[1:-len(tokenizer.sot_sequence)])

	may_flush_segment(curr_tokens)

	# Get the index of the <\|startoftranscript\|> tokens (to get proba of silence later)
	if is_sot(curr_tokens):
	has_started = len(curr_tokens) > 1 or not model.is_multilingual
	if no_speech_threshold is not None:
	sot_index = curr_tokens.tolist().index(tokenizer.sot)
	else:
	sot_index = None

	# Keep the last token only
	if has_started:
	segment_tokens[-1].append(curr_tokens[-1].item())

	# Accumulate tokens
	if has_started:
	chunk_tokens.append(curr_tokens)
	if not is_sot(curr_tokens):
	chunk_tokens_nosot.append(curr_tokens[-1].item())
	else:
	if verbose and not whisper_options["verbose"]:
	# Reproduce whisper verbose (1/2)
	print("Detecting language using up to the first 30 seconds. Use `--language` to specify the language")

	embedding_weights = None
	def hook_output_logits(layer, ins, outs):
	nonlocal no_speech_prob, chunk_logprobs, segment_tokens, chunk_tokens, chunk_tokens_nosot, last_chunk_token, embedding_weights, has_started, language, language_probs

	if embedding_weights is None:
	embedding_weights = torch.transpose(model.decoder.token_embedding.weight, 0, 1).to(outs[0].dtype)

	# Get the probability of silence
	if sot_index is not None and no_speech_prob is None:
	logits = (outs[0][sot_index,:] @ embedding_weights).float()
	logits = logits.softmax(dim=-1)
	no_speech_prob = logits[tokenizer.no_speech].item()

	# Get language probabilities
	if language is None and sot_index is not None and model.is_multilingual:
	index_start = tokenizer.sot + 1
	index_end = index_start + len(tokenizer.all_language_tokens)
	logits = (outs[0][sot_index,:] @ embedding_weights).float()
	language_probs = logits[index_start:index_end].softmax(dim=-1)
	language_probs = dict(zip(whisper.tokenizer.LANGUAGES, language_probs.tolist()))

	# Get the log-probabilities of tokens (we don't know yet which one will be chosen)
	if has_started:
	logits = (outs[0][-1:,:] @ embedding_weights).float()
	tokens = torch.cat(chunk_tokens).unsqueeze(0)
	for logit_filter in logit_filters:
	logit_filter.apply(logits, tokens)
	logits = F.log_softmax(logits.squeeze(0), dim=-1)
	chunk_logprobs.append(logits)

	if WHIPSER_GE_20230306 and has_reached_decoding_limit():
	last_chunk_token = torch.argmax(logits).item()
	else:
	last_chunk_token = None

	try:

	# Add hooks to the model, to get tokens and attention weights on the fly
	all_hooks = []
	all_hooks.append(model.encoder.conv1.register_forward_hook(hook_mfcc))
	all_hooks.append(model.decoder.token_embedding.register_forward_hook(hook_input_tokens))
	nblocks = len(model.decoder.blocks)
	j = 0
	for i, block in enumerate(model.decoder.blocks):
	if i < nblocks - word_alignement_most_top_layers:
	continue
	all_hooks.append(
	block.cross_attn.register_forward_hook(
	lambda layer, ins, outs, index=j: hook_attention_weights(layer, ins, outs, index))
	)
	j += 1
	if compute_word_confidence or no_speech_threshold is not None:
	all_hooks.append(model.decoder.ln.register_forward_hook(hook_output_logits))

	transcription = model.transcribe(audio, **whisper_options)

	finally:

	# Remove hooks
	for hook in all_hooks:
	hook.remove()

	# Finalize (collect last segment)
	may_flush_segment()
	segment_tokens.pop(-1)

	token_special_idx = min(tokenizer.sot, tokenizer.eot)
	def filter_tokens(tokens):
	while len(tokens) and tokens[0] >= token_special_idx:
	tokens = tokens[1:]
	while len(tokens) and tokens[-1] >= token_special_idx:
	tokens = tokens[:-1]
	return tokens

	assert len(segment_tokens) == len(timestamped_word_segments), f"Inconsistent number of segments: tokens ({len(segment_tokens)}) != timestamped_word_segments ({len(timestamped_word_segments)})"
	assert len(segment_avglogprobs) == len(segment_tokens), f"Inconsistent number of segments: avg logprobs ({len(segment_avglogprobs)}) != tokens ({len(segment_tokens)})"
	assert len(segment_logprobs) == len(segment_tokens), f"Inconsistent number of segments: logprobs ({len(segment_logprobs)}) != tokens ({len(segment_tokens)})"

	whisper_segments = transcription["segments"]
	l1 = len(whisper_segments)
	l2 = len(timestamped_word_segments)
	if l1 != l2 and l1 != 0:
	logger.warning(f"Inconsistent number of segments: whisper_segments ({l1}) != timestamped_word_segments ({l2})")
	assert l1 == l2 or l1 == 0, f"Inconsistent number of segments: whisper_segments ({l1}) != timestamped_word_segments ({l2})"

	logger.debug("Compile results")
	words = []
	for i, (segment, timestamped_words, token, avglogprob, logprobs) in enumerate(zip(whisper_segments, timestamped_word_segments, segment_tokens, segment_avglogprobs, segment_logprobs)):
	timestamped_tokens = filter_tokens(token)
	whisper_tokens = filter_tokens(segment["tokens"])
	if timestamped_tokens != whisper_tokens:
	if len(timestamped_tokens) == len(whisper_tokens) + 1:
	logger.warning(f"An additional token was added on segment {i}")
	elif WHIPSER_GE_20230306 and len(whisper_tokens) == 0:
	logger.warning(f"Whisper has empty segment {i}")
	assert segment["end"] == segment["start"], f"Fatal Error: Got empty segment {i} with non-zero duration"
	segment["tokens"] = timestamped_tokens
	segment["text"] = tokenizer.decode(timestamped_tokens)
	else:
	assert len(timestamped_tokens) < len(whisper_tokens) and timestamped_tokens == whisper_tokens[:len(timestamped_tokens)], \
	f"Fatal Error: Got inconsistent text for segment {i}:\n({len(timestamped_tokens)})\n{tokenizer.decode_with_timestamps(timestamped_tokens)}\n{timestamped_tokens}\n!=\n({len(whisper_tokens)})\n{tokenizer.decode_with_timestamps(whisper_tokens)}\n{whisper_tokens[:len(timestamped_tokens)]}"
	segment["tokens"] = token if WHIPSER_GE_20230306 else timestamped_tokens # tokens include special timestamp tokens since 20230306
	segment["text"] = tokenizer.decode(segment["tokens"])
	logger.warning(f"Text had to be shortned on segment {i}:\n{tokenizer.decode(timestamped_tokens)}\n!=\n{tokenizer.decode(whisper_tokens)}")
	timestamped_words[-1]["avg_logprob_reliable"] = False

	offset = segment["seek"] * HOP_LENGTH / SAMPLE_RATE
	for timestamped_word in timestamped_words:
	timestamped_word["start"] += offset
	timestamped_word["end"] += offset
	timestamped_word["idx_segment"] = i

	if compute_word_confidence:
	if "avg_logprob_reliable" not in timestamped_words[-1] or timestamped_words[-1]["avg_logprob_reliable"]:
	# assert abs(segment["avg_logprob"] - avglogprob) < 1e-2, f"Fatal Error: Got inconsistent logprob for segment {i}: {segment['avg_logprob']} != {avglogprob}"
	if abs(segment["avg_logprob"] - avglogprob) >= 1e-2:
	logger.warning(f"Recomputed different logprob for segment {i}: {avglogprob} != {segment['avg_logprob']}")
	if include_punctuation_in_confidence:
	segment["confidence"] = round_confidence(logprobs.mean().exp().item())
	else:
	logprobs_nopunc = []
	i_end = 0
	for timestamped_word in timestamped_words:
	i_start = i_end
	tokens = timestamped_word["tokens"]
	i_end += len(tokens)

	assert i_end <= len(logprobs), f"Fatal Error: Got out-of-bound index for segment {i}: {i_end} > {len(logprobs)}"
	if include_punctuation_in_confidence:
	word_logprobs = logprobs[i_start:i_end]
	else:
	while len(tokens) > 1 and len(tokens[-1]) and tokens[-1][-1] in _punctuation: # Note: look at the last character of token, to take into account "...", "!!", etc.
	tokens = tokens[:-1]
	word_logprobs = logprobs[i_start:i_start + len(tokens)]
	logprobs_nopunc.append(word_logprobs)

	timestamped_word["confidence"] = round_confidence(word_logprobs.mean().exp().item() if len(word_logprobs) else 0.0)

	if i_end not in [len(logprobs), len(logprobs)-1]:
	logger.warning(f"Got inconsistent length for segment {i} ({len(logprobs)} != {i_end}). Some words have been ignored.")
	if not include_punctuation_in_confidence:
	logprobs_nopunc = torch.cat(logprobs_nopunc)
	segment["confidence"] = round_confidence(logprobs_nopunc.mean().exp().item())

	words.extend(timestamped_words)

	if language_probs:
	transcription["language_probs"] = language_probs

	return transcription, words

	def _transcribe_timestamped_naive(
	model,
	audio,
	remove_punctuation_from_words,
	compute_word_confidence,
	include_punctuation_in_confidence,
	refine_whisper_precision_nframes,
	alignment_heads,
	plot_word_alignment,
	word_alignement_most_top_layers,
	detect_disfluencies,
	trust_whisper_timestamps,
	min_word_duration,
	**whisper_options,
	):
	verbose = whisper_options["verbose"]
	whisper_options["verbose"] = None if whisper_options["verbose"] is True else whisper_options["verbose"] # We will print intermediate results ourselves
	language = whisper_options["language"]
	refine_whisper_precision_sec = refine_whisper_precision_nframes * AUDIO_TIME_PER_TOKEN

	word_alignement_most_top_layers = float("inf") if word_alignement_most_top_layers is None else word_alignement_most_top_layers

	audio = get_audio_tensor(audio)
	audio_duration = audio.shape[-1] / SAMPLE_RATE

	if verbose and language is None and not whisper_options["verbose"]:
	# Reproduce whisper verbose (1/2)
	print("Detecting language using up to the first 30 seconds. Use `--language` to specify the language")

	tokenizer = get_tokenizer(model, task=whisper_options["task"], language=language)

	language_probs = None
	def hook_output_logits(layer, ins, outs):
	nonlocal language_probs, tokenizer

	# Get language probabilities
	if language_probs is None:
	if outs.shape[1] == 1:
	embedding_weights = torch.transpose(model.decoder.token_embedding.weight, 0, 1).to(outs[0].dtype)
	index_start = tokenizer.sot + 1
	index_end = index_start + len(tokenizer.all_language_tokens)
	logits = (outs[0][0,:] @ embedding_weights).float()
	language_probs = logits[index_start:index_end].softmax(dim=-1)
	language_probs = dict(zip(whisper.tokenizer.LANGUAGES, language_probs.tolist()))
	else:
	language_probs = False

	all_hooks = []
	if model.is_multilingual:
	all_hooks.append(model.decoder.ln.register_forward_hook(hook_output_logits))

	try:
	transcription = model.transcribe(audio, **whisper_options)
	finally:
	for hook in all_hooks:
	hook.remove()

	if verbose and language is None and not whisper_options["verbose"]:
	# Reproduce whisper verbose (2/2)
	print(f"Detected language: {whisper.tokenizer.LANGUAGES[transcription['language']].title()}")
	sys.stdout.flush()

	language = norm_language(transcription["language"])
	use_space = should_use_space(language)

	n_mels = model.dims.n_mels if hasattr(model.dims, "n_mels") else 80

	attention_weights = [[] for _ in range(min(word_alignement_most_top_layers,len(model.decoder.blocks)))]

	try:

	all_hooks = []

	# Hook the model
	nblocks = len(model.decoder.blocks)
	j = 0
	for i, block in enumerate(model.decoder.blocks):
	if i < nblocks - word_alignement_most_top_layers:
	continue
	all_hooks.append(
	block.cross_attn.register_forward_hook(
	lambda layer, ins, outs, index=j: attention_weights.__setitem__(index, outs[-1])
	)
	)
	j += 1


	# When not relying on Whisper timestamps
	current_tokens = []
	token_to_idx_segment = []

	words = []
	previous_end = 0
	whisper_segments = transcription["segments"]
	for i_segment, segment in enumerate(whisper_segments):

	# Note: this could also be a fix to issue #61 where a "<\|te\|>" token was predicted
	# segment["tokens"] = [t for t in segment["tokens"] if t < tokenizer.eot or t >= tokenizer.timestamp_begin]

	start = end = tokens = None
	if trust_whisper_timestamps:

	start = segment["start"]
	end = segment["end"]
	if end < start:
	# Whisper is wrong on the prediction of segment end
	end = min(audio_duration, start + SEGMENT_DURATION)

	start_margin_min = start - refine_whisper_precision_sec
	start_margin_max = start + refine_whisper_precision_sec
	if start >= audio_duration - min_word_duration or (previous_end >= start_margin_min and previous_end <= start_margin_max):
	# Make start as accurate as possible (as the decoding will start with timestamp <\|0\|>)
	start = previous_end
	else:
	# Fallback
	start = start_margin_min

	if start > audio_duration - min_word_duration:
	# Skip last segment if too short
	logger.warning(f"Skipping segment outside of audio duration {audio_duration} (original: {segment['start']}-{segment['end']}, new: {start}-XXX)")
	continue

	end_margin_min = end - refine_whisper_precision_sec
	end_margin_max = end + refine_whisper_precision_sec
	if i_segment < len(whisper_segments) - 1:
	# Try to enforce:
	# end + min_word_duration <= next start + refine_whisper_precision_sec
	end_margin_max2 = whisper_segments[i_segment + 1]["start"] + refine_whisper_precision_sec - min_word_duration
	if end_margin_max2 >= end_margin_min:
	end_margin_max = min(end_margin_max2, end_margin_max)
	end = min(audio_duration, end_margin_max)

	if end < start + min_word_duration:
	logger.warning(f"Got super short segment (original from whisper: {segment['start']}-{segment['end']}, new: {start, end})")
	end = min(audio_duration, start + min_word_duration)
	if end <= start:
	logger.warning(f"Skipping this short segment occuring too close to the end of the audio")
	continue

	tokens = segment["tokens"]

	else:

	seek = segment["seek"]
	new_tokens = segment["tokens"]
	if not len(new_tokens):
	continue
	# Add timestamps that will be needed after
	if new_tokens[0] < tokenizer.timestamp_begin:
	relative_start = segment["start"] - (seek * HOP_LENGTH / SAMPLE_RATE)
	start_token = round(relative_start * SAMPLE_RATE / AUDIO_SAMPLES_PER_TOKEN) + tokenizer.timestamp_begin
	new_tokens = [start_token] + new_tokens
	if new_tokens[-1] < tokenizer.timestamp_begin:
	relative_end = segment["end"] - (seek * HOP_LENGTH / SAMPLE_RATE)
	end_token = round(relative_end * SAMPLE_RATE / AUDIO_SAMPLES_PER_TOKEN) + tokenizer.timestamp_begin
	new_tokens = new_tokens + [end_token]

	current_tokens.extend(new_tokens)
	token_to_idx_segment.extend([i_segment] * len(new_tokens))

	next_seek = whisper_segments[i_segment+1]["seek"] if i_segment < len(whisper_segments) - 1 else None
	if seek != next_seek:
	start = float(seek * HOP_LENGTH / SAMPLE_RATE)
	assert start < audio_duration, f"Got start {start} which is outside of audio duration {audio_duration}"
	end = min(start + SEGMENT_DURATION, audio_duration)
	tokens = current_tokens

	if tokens is None or not len(tokens):
	continue

	start_sample = min(round(start * SAMPLE_RATE), audio.shape[-1])
	end_sample = min(round(end * SAMPLE_RATE), audio.shape[-1])

	# Extract features on the audio segment
	sub_audio = audio_minimum_padding(audio[start_sample:end_sample])

	mfcc = whisper.log_mel_spectrogram(sub_audio, n_mels).to(model.device)
	mfcc = whisper.pad_or_trim(mfcc, N_FRAMES)
	mfcc = mfcc.unsqueeze(0)

	segment_tokens_check = []
	if tokens[0] >= tokenizer.timestamp_begin:
	segment_tokens_check.append(tokens[0])
	while tokens[0] >= tokenizer.timestamp_begin:
	tokens = tokens[1:]
	assert len(tokens), "Got transcription with only timestamps!"
	last_token_check = None
	while tokens[-1] >= tokenizer.timestamp_begin:
	last_token_check = tokens[-1]
	tokens = tokens[:-1]

	tokens = [
	*tokenizer.sot_sequence,
	tokenizer.timestamp_begin,
	] + tokens

	i_start = len(tokenizer.sot_sequence)

	with torch.no_grad():
	logprobs = model(mfcc, torch.Tensor(tokens).int().to(model.device).unsqueeze(0))
	logprobs = F.log_softmax(logprobs, dim=-1)

	end_token = tokenizer.timestamp_begin + round(min(N_FRAMES * HOP_LENGTH, end_sample - start_sample) // AUDIO_SAMPLES_PER_TOKEN)
	tokens = tokens[i_start:] + [end_token]
	attention_weights = [w[:, :, i_start-1:, :] for w in attention_weights]

	ws = perform_word_alignment(
	tokens,
	attention_weights,
	tokenizer,
	use_space=use_space,
	alignment_heads=alignment_heads,
	remove_punctuation_from_words=remove_punctuation_from_words,
	refine_whisper_precision_nframes=refine_whisper_precision_nframes,
	detect_disfluencies=detect_disfluencies,
	mfcc=mfcc,
	plot=plot_word_alignment,
	)

	segment_logprobs = []
	i_token = 1

	for word in ws:

	word["start"] = round(word["start"] + start, 2)
	word["end"] = round(word["end"] + start, 2)

	if trust_whisper_timestamps:
	word.update({"idx_segment": i_segment})
	else:
	assert i_token < len(tokens)
	assert not len(word["tokens_indices"]) or word["tokens_indices"][0] == tokens[i_token]
	word.update({"idx_segment": token_to_idx_segment[i_token]})
	i_token += len(word["tokens"])
	while i_token < len(tokens) and tokens[i_token] >= tokenizer.timestamp_begin:
	i_token += 1

	tok_indices = word["tokens_indices"]
	segment_tokens_check.extend(tok_indices)

	if compute_word_confidence:
	tok = word["tokens"]
	i_end = i_start + len(tok)
	if include_punctuation_in_confidence:
	while len(tok) > 1 and len(tok[-1]) and tok[-1][-1] in _punctuation: # Note: look at the last character of token, to take into account "...", "!!", etc.
	tok = tok[:-1]
	tok_indices = tok_indices[:-1]
	word_logprobs = [logprobs[:, step, tok] for (step, tok) in zip(range(i_start, i_start + len(tok_indices)), tok_indices)]
	i_start = i_end
	if len(word_logprobs):
	word_logprobs = torch.cat(word_logprobs)
	segment_logprobs.append(word_logprobs)
	word_confidence = word_logprobs.mean().exp().item()
	else:
	word_confidence = 0
	word.update({"confidence": round_confidence(word_confidence)})

	words.append(word)

	if verbose:
	print_timestamped(word)

	if last_token_check is not None:
	segment_tokens_check.append(last_token_check)
	if trust_whisper_timestamps:
	if segment_tokens_check != segment["tokens"]:
	assert len(segment_tokens_check) < len(segment["tokens"]) and segment_tokens_check[:-1] == segment["tokens"][:len(segment_tokens_check)-1], \
	f"Got inconsistent tokens: {tokenizer.decode(segment_tokens_check)} != {tokenizer.decode(segment['tokens'])}"
	segment["tokens"] = segment_tokens_check
	segment["text"] = tokenizer.decode(segment["tokens"])
	# else: TODO

	if len(segment_logprobs):
	segment.update({"confidence": round_confidence(torch.cat(segment_logprobs).mean().exp().item())})

	if len(ws):
	previous_end = ws[-1]["end"]

	if not trust_whisper_timestamps:
	current_tokens = []
	token_to_idx_segment = []

	finally:

	# Remove hooks
	for hook in all_hooks:
	hook.remove()

	if language_probs:
	transcription["language_probs"] = language_probs

	return (transcription, words)

	def get_audio_tensor(audio, device="cpu"):
	if isinstance(audio, str):
	audio = whisper.load_audio(audio)
	if isinstance(audio, np.ndarray):
	audio = torch.Tensor(audio)
	else:
	assert isinstance(audio, torch.Tensor), f"Got unexpected audio of type {type(audio)}"
	return audio.to(device)

	def audio_minimum_padding(audio):
	if audio.shape[-1] <= 200:
	return whisper.pad_or_trim(audio, 201)
	return audio


	def should_use_space(language):
	return norm_language(language) not in ["zh", "ja", "th", "lo", "my", "yue"]

	def norm_language(language):
	if language is None:
	return "en"
	return whisper.tokenizer.TO_LANGUAGE_CODE.get(language.lower(), language)

	def print_timestamped(w):
	line = f"[{format_timestamp(w['start'])} --> {format_timestamp(w['end'])}] {w['text']}\n"
	# compared to just `print(line)`, this replaces any character not representable using
	# the system default encoding with an '?', avoiding UnicodeEncodeError.
	sys.stdout.write(line.encode(sys.getdefaultencoding(), errors="replace").decode())
	sys.stdout.flush()


	def get_logit_filters(model, whisper_options, prompt = None):
	decoding_options = get_decoding_options(whisper_options)
	if "initial_prompt" in decoding_options:
	prompt0 = decoding_options.pop("initial_prompt")
	if prompt is None:
	prompt = prompt0
	if prompt is not None:
	decoding_options["prompt"] = prompt
	decoding_options = whisper.DecodingOptions(
	without_timestamps=False,
	max_initial_timestamp=1.0,
	prefix=None,
	suppress_blank=True,
	**decoding_options
	)

	# This performs some checks on the options
	decoding_task = whisper.decoding.DecodingTask(model, decoding_options)
	return decoding_task.logit_filters

	def get_decoding_options(whisper_options):
	return dict([(k,v) for (k,v) in whisper_options.items()
	if k not in [
	"no_speech_threshold",
	"logprob_threshold",
	"compression_ratio_threshold",
	"condition_on_previous_text",
	"verbose",
	]
	])

	def get_tokenizer(model, task="transcribe", language="en"):
	try:
	return whisper.tokenizer.get_tokenizer(
	model.is_multilingual,
	num_languages=model.num_languages if hasattr(model, "num_languages") else 99,
	task=task, language=language
	)
	except TypeError: # Old openai-whisper version
	return whisper.tokenizer.get_tokenizer(
	model.is_multilingual,
	task=task, language=language
	)

	def perform_word_alignment(
	tokens,
	attention_weights,
	tokenizer,
	use_space=True,
	mfcc=None,
	refine_whisper_precision_nframes=0,
	remove_punctuation_from_words=False,
	include_punctuation_in_timing=False, # Was True before 1.9
	unfinished_decoding=False,
	alignment_heads=None,
	medfilt_width=9,
	qk_scale=1.0,
	detect_disfluencies=True,
	subwords_can_be_empty=True, # Was False before 1.11
	plot=False,
	debug=False,
	):
	"""
	Perform word alignment on the given tokens and attention weights.
	Returns a list of (word, start_time, end_time) tuples.

	tokens: list of tokens (integers)
	attention_weights: list of attention weights (torch tensors)
	tokenizer: tokenizer used to tokenize the text
	use_space: whether to use spaces to split the tokens into words (should be true for all languages except Japanese, Chinese, ...)
	mfcc: MFCC features (used to identify padded region, and for plotting)
	refine_whisper_precision_nframes: precision time
	remove_punctuation_from_words: whether to remove punctuation from words
	include_punctuation_in_timing: whether to include punctuation in the timing of (previous) words
	unfinished_decoding: whether the decoding is unfinished (e.g. because the model is stuck)
	alignment_heads: list of attention heads to use for alignment
	medfilt_width: width of the median filter used to smooth the attention weights
	qk_scale: scale factor applied to the attention weights
	plot: whether to plot the word alignment
	debug: whether to print debug information
	"""

	assert len(tokens) > 1, f"Got unexpected sequence of tokens of length {len(tokens)} {tokenizer.decode_with_timestamps(tokens)}"
	start_token = tokens[0] - tokenizer.timestamp_begin
	end_token = tokens[-1] - tokenizer.timestamp_begin

	# Check start / end tokens
	if start_token < 0:
	raise RuntimeError(f"Missing start token in: {tokenizer.decode_with_timestamps(tokens)}")
	if len(tokens) == 1 or end_token < 0:
	# This can happens when Whisper is stucked as a Language Model
	if debug:
	logger.debug(f"Missing end token in {tokenizer.decode_with_timestamps(tokens)}")
	end_token = N_FRAMES // 2
	if end_token == start_token and refine_whisper_precision_nframes == 0:
	if debug:
	logger.debug(f"Got empty segment in {tokenizer.decode_with_timestamps(tokens)}")
	return []

	# Let a minimal duration given the number of tokens (see https://github.com/linto-ai/whisper-timestamped/issues/67)
	end_token = min(N_FRAMES // 2, max(end_token, start_token + len(tokens)))

	# Put some margin around the segment
	if refine_whisper_precision_nframes > 0:
	start_token = max(start_token - refine_whisper_precision_nframes, 0)
	end_token = min(end_token + refine_whisper_precision_nframes, N_FRAMES // 2)

	if end_token <= start_token:
	raise RuntimeError(f"Got segment with null or negative duration {tokenizer.decode_with_timestamps(tokens)}: {start_token} {end_token}")

	start_time = start_token * AUDIO_TIME_PER_TOKEN
	# end_time = end_token * AUDIO_TIME_PER_TOKEN

	split_tokens = split_tokens_on_spaces if use_space else split_tokens_on_unicode
	words, word_tokens, word_tokens_indices = split_tokens(tokens, tokenizer, remove_punctuation_from_words=remove_punctuation_from_words)

	# If the last token is a punctuation that comes after a word
	# group this final punctuation with the final timestamp
	# This is to avoid assigning the final punctuation to a big silence or a noise/music background coming after
	num_punctuations_per_tokens = [
	0 if len(w) == 1 or w[-1] not in _punctuation else 1
	for w in word_tokens
	]
	if include_punctuation_in_timing:
	num_punctuations_per_tokens[:-2]=[0]*(len(num_punctuations_per_tokens)-2)

	for i, w in enumerate(attention_weights):
	assert w.shape[-2] == len(tokens), f"Attention weights have wrong shape: {w.shape[-2]} (expected {len(tokens)})."
	weights = torch.cat(attention_weights) # layers * heads * tokens * frames

	num_tokens = weights.shape[-2]
	num_frames = end_token - start_token
	if num_tokens > num_frames:
	logger.warning(f"Too much text ({num_tokens} tokens) for the given number of frames ({num_frames}) in: {tokenizer.decode_with_timestamps(tokens)}\nThe end of the text will be removed.")
	return perform_word_alignment(
	tokens[:num_frames-1] + [tokens[-1]],
	[torch.cat([w[:, :, :num_frames-1, :], w[:, :, -1:, :]], dim=-2)
	for w in attention_weights],
	tokenizer,
	use_space=use_space,
	refine_whisper_precision_nframes=refine_whisper_precision_nframes,
	medfilt_width=medfilt_width,
	qk_scale=qk_scale,
	alignment_heads=alignment_heads,
	mfcc=mfcc,
	plot=plot,
	remove_punctuation_from_words=remove_punctuation_from_words,
	detect_disfluencies=detect_disfluencies,
	subwords_can_be_empty=subwords_can_be_empty,
	unfinished_decoding=True,
	debug=debug,
	)

	assert end_token <= weights.shape[-1]
	assert len(tokens) == num_tokens

	weights = weights[..., start_token: end_token].cpu() # layers * heads * tokens * frames

	if alignment_heads is None:
	weights = weights.reshape(-1, weights.shape[-2:]) # N tokens * frames
	else:
	weights = torch.stack([weights[l][h] for l, h in alignment_heads.indices().T])
	weights = median_filter(weights, (1, 1, medfilt_width))
	weights = torch.tensor(weights * qk_scale).softmax(dim=-1)
	weights = weights.mean(axis=(0)) # average over layers and heads # tokens * frames
	weights = weights / weights.norm(dim=-2, keepdim=True) # This was before the mean before 1.9
	weights = -weights.double().numpy()
	worse_weight = 0

	# Get the limit of audio duration
	max_duration = None
	if mfcc is not None:
	max_duration = find_start_padding(mfcc)
	if max_duration is not None:
	max_duration = max_duration // 2

	# Enforce the max duration
	if max_duration:
	if start_token >= max_duration:
	logger.warning(f"Got start time outside of audio boundary")
	else:
	weights[:-1, max_duration:] = worse_weight

	# Encourage to start early
	weights[0, 0] = weights.min()
	# weights[0, refine_whisper_precision_nframes*2:] = worse_weight

	if subwords_can_be_empty:
	step_pattern = dtw.stepPattern.symmetric1
	else:
	# Similar as "symmetric1" but without the possibility to have the same timestamp for two tokens
	step_pattern = dtw.stepPattern.StepPattern(dtw.stepPattern._c(
	1, 1, 1, -1,
	1, 0, 0, 1,
	2, 0, 1, -1,
	2, 0, 0, 1,
	))
	alignment = dtw.dtw(weights, step_pattern=step_pattern)

	global num_alignment_for_plot
	num_alignment_for_plot += 1

	if plot:
	import matplotlib.pyplot as plt
	import matplotlib.ticker as ticker

	plot_mfcc = 1 if mfcc is not None else 0
	plot_disfluencies = 1 if detect_disfluencies else 0
	nplots = (1 + plot_mfcc + plot_disfluencies)

	plt.subplots(nplots, 1, figsize=(16, 9), gridspec_kw={'height_ratios': [3] + [1] * (nplots - 1)})
	plt.subplot(nplots, 1, 1, frameon=False)

	plt.imshow(-weights, aspect="auto")
	plt.plot(alignment.index2s, alignment.index1s, color="red")

	xticks = np.arange(0, weights.shape[1], 1 / AUDIO_TIME_PER_TOKEN)
	xticklabels = [round_timestamp(x) for x in xticks * AUDIO_TIME_PER_TOKEN + start_time]

	ylims = plt.gca().get_ylim()

	ax = plt.gca()
	ax.tick_params('both', length=0, width=0, which='minor', pad=6)

	ax.yaxis.set_ticks_position("left")
	ax.yaxis.set_label_position("left")
	ax.invert_yaxis()
	ax.set_ylim(ylims)

	major_ticks = [-0.5]
	minor_ticks = []
	current_y = 0

	for word, word_token in zip(words, word_tokens):
	minor_ticks.append(current_y + len(word_token) / 2 - 0.5)
	current_y += len(word_token)
	major_ticks.append(current_y - 0.5)

	words_with_subwords = ["\|".join(s).strip() for (w, s) in zip(words, word_tokens)]

	ax.yaxis.set_minor_locator(ticker.FixedLocator(minor_ticks))
	ax.yaxis.set_minor_formatter(
	ticker.FixedFormatter(words_with_subwords))
	ax.set_yticks(major_ticks)
	ax.yaxis.set_major_formatter(ticker.NullFormatter())
	for y in major_ticks:
	plt.axhline(y, color="black", linestyle="dashed")

	plt.ylabel("Words")

	if plot_mfcc:
	plt.xticks(xticks)
	plt.setp(plt.gca().get_xticklabels(), visible=False)

	xticks *= 2

	plt.subplot(nplots, 1, 2, frameon=False)
	plt.imshow(mfcc[0, :, start_token * 2: end_token * 2].cpu(), aspect="auto", origin="lower")
	plt.yticks([])
	plt.ylabel("MFCC")

	plt.xticks(xticks, xticklabels)
	plt.xlabel("Time (s)")

	jumps = np.diff(alignment.index1s)
	jumps = np.pad(jumps, (1, 0), constant_values=1)
	jumps = jumps.astype(bool)
	jumps = alignment.index2s[jumps]
	jumps = np.pad(jumps, (0, 1), constant_values=alignment.index2s[-1])

	jumps_start = jumps
	disfluences = {}
	if detect_disfluencies:
	jumps_start = copy.copy(jumps)

	for (i_token, (tok, begin, end)) in enumerate(zip(tokens, jumps[:-1], jumps[1:])):

	# Find local maxima in the portion of attention weights
	attention_weights = -weights[i_token, begin:end]
	peaks, properties = find_peaks(attention_weights,
	width=3,
	prominence=0.02,
	)
	# If more than
	if len(peaks) > 1:
	if "left_ips" in properties:
	left = [round(x) for x in properties["left_ips"]]
	else:
	left = properties["left_bases"]

	new_begin = left[-1] + begin

	jumps_start[i_token] = new_begin

	if new_begin != begin:
	is_punctuation = tokenizer.decode_with_timestamps([tok]) in _punctuation
	if not is_punctuation:
	disfluences[i_token] = (begin, jumps_start[i_token])
	else:
	disfluences[i_token+1] = (begin, end)

	if plot:
	plt.subplot(nplots, 1, 2 + plot_mfcc, frameon=False)
	plt.plot(range(begin,end), attention_weights)
	plt.xlim(0, end)

	for i, p in enumerate(peaks):
	color = 'red' if (len(peaks)>1 and i<len(peaks)-1) else 'green'
	plt.vlines(begin+p, 0, 1, color=color, linestyle="--")

	if "left_bases" in properties:
	def barxxy(start, end, y, **kwargs):
	middle = (start + end) / 2
	plt.bar(middle, y, width=end-start, **kwargs)
	color = 'red' if len(peaks)>1 else 'green'
	barxxy(begin+properties["left_bases"], begin+properties["right_bases"], properties.get("prominences",[1]*len(properties["left_bases"])), alpha=0.5,
	# put a line with a custom color
	linewidth=1, edgecolor=color
	)
	if "left_ips" in properties:
	for left in properties["left_ips"]:
	plt.vlines(begin+left, 0, 0.5, color='green', linestyle=':')
	for right in properties["right_ips"]:
	plt.vlines(begin+right, 0, 0.5, color='red', linestyle=':')


	# display the word-level timestamps in a table
	word_boundaries = np.cumsum([len(t) for t in word_tokens])
	word_boundaries = np.pad(word_boundaries, (1, 0))
	begin_times = jumps_start[word_boundaries[:-1]]
	end_times = jumps[word_boundaries[1:] - num_punctuations_per_tokens]

	begin_times = begin_times * AUDIO_TIME_PER_TOKEN
	end_times = end_times * AUDIO_TIME_PER_TOKEN

	if detect_disfluencies:
	to_be_added = []
	i_start = 0
	for i_word, toks in enumerate(word_tokens[:-1]):
	i_end = i_start + len(toks)
	if i_start in disfluences and i_word > 0:
	begin, end = disfluences[i_start]
	begin *= AUDIO_TIME_PER_TOKEN
	end *= AUDIO_TIME_PER_TOKEN
	to_be_added.append((i_word, begin, end))
	i_start = i_end
	# Add from the end to avoid messing up the indices
	for (i_word, begin, end) in to_be_added[-1::-1]:
	words.insert(i_word, DISFLUENCY_MARK)
	word_tokens.insert(i_word, [])
	word_tokens_indices.insert(i_word, [])
	begin_times = np.insert(begin_times, i_word, begin)
	end_times = np.insert(end_times, i_word, end)

	# Ignore start / end tokens
	if not refine_whisper_precision_nframes:
	begin_times[1] = begin_times[0]
	if not refine_whisper_precision_nframes:
	end_times[-2] = end_times[-1]
	if unfinished_decoding:
	words = words[1:]
	word_tokens = word_tokens[1:]
	word_tokens_indices = word_tokens_indices[1:]
	begin_times = begin_times[1:]
	end_times = end_times[1:]
	else:
	words = words[1:-1]
	word_tokens = word_tokens[1:-1]
	word_tokens_indices = word_tokens_indices[1:-1]
	begin_times = begin_times[1:-1]
	end_times = end_times[1:-1]

	if plot:
	ymin = 1

	plt.subplot(nplots, 1, 1)
	for i, (w, ws, begin, end) in enumerate(zip(words, word_tokens, begin_times, end_times)):
	ymax = ymin + len(ws)
	if mfcc is None:
	plt.text(begin / AUDIO_TIME_PER_TOKEN, num_tokens-0.5, w, ha="left", va="top", color="red")
	for x in [begin, end,]:
	plt.axvline(x / AUDIO_TIME_PER_TOKEN, color="red", linestyle="dotted",
	ymin=1-ymin/num_tokens,
	ymax=0, # 1-ymax/num_tokens,
	)
	ymin = ymax

	if plot_mfcc:
	plt.subplot(nplots, 1, 2)
	for i, (w, begin, end) in enumerate(zip(words, begin_times, end_times)):
	plt.text(begin * 2 / AUDIO_TIME_PER_TOKEN, mfcc.shape[-2]*1.05, w, ha="left", va="bottom", color="red")
	for x in [begin, end,]:
	plt.axvline(x * 2 / AUDIO_TIME_PER_TOKEN, color="red", linestyle="dotted")

	if isinstance(plot, str):
	plt.savefig(f"{plot}.alignment{num_alignment_for_plot:03d}.jpg", bbox_inches='tight', pad_inches=0)
	else:
	plt.show()

	return [
	dict(
	text=word,
	start=round_timestamp(begin + start_time),
	end=round_timestamp(end + start_time),
	tokens=tokens,
	tokens_indices=tokens_indices,
	)
	for word, begin, end, tokens, tokens_indices in zip(words, begin_times, end_times, word_tokens, word_tokens_indices)
	if not word.startswith("<\|")
	]

	def find_start_padding(mfcc):
	""" Return start of padding given the mfcc, or None if there is no padding """
	last_mfcc = mfcc[0, :, -1]
	if torch.min(last_mfcc) == torch.max(last_mfcc) == 0:
	candidate_index = mfcc.shape[-1] - 2
	while candidate_index > 0:
	candidate = mfcc[0, :, candidate_index]
	if not torch.equal(candidate, last_mfcc):
	return candidate_index + 1
	candidate_index -= 1
	return 0 # WTF!?

	def round_confidence(x):
	return round(x, 3)

	def round_timestamp(x):
	return round(x, 2)

	_punctuation = "".join(c for c in string.punctuation if c not in ["-", "'"]) + "。，！？：”、…"

	def split_tokens_on_unicode(tokens: list, tokenizer, remove_punctuation_from_words=False, isolate_punctuations=False):
	words = []
	word_tokens = []
	word_tokens_indices = []
	current_tokens = []

	for token in tokens:
	current_tokens.append(token)
	decoded = tokenizer.decode_with_timestamps([t for t in current_tokens if t < tokenizer.eot or t >= tokenizer.timestamp_begin])
	if "\ufffd" not in decoded:
	empty_tokens = [""] * (len(current_tokens)-1)
	punctuation = not isolate_punctuations and (decoded.strip() and decoded.strip() in _punctuation)
	previous_special = len(word_tokens_indices) > 0 and (word_tokens_indices[-1][-1] >= tokenizer.timestamp_begin)
	if punctuation and not previous_special:
	if len(words) == 0:
	words = [""]
	word_tokens = [[]]
	if not remove_punctuation_from_words:
	words[-1] += decoded
	word_tokens[-1].extend(empty_tokens + [decoded])
	word_tokens_indices[-1].extend(current_tokens)
	else:
	words.append(decoded)
	word_tokens.append(empty_tokens + [decoded])
	word_tokens_indices.append(current_tokens)
	current_tokens = []

	return words, word_tokens, word_tokens_indices


	def split_tokens_on_spaces(tokens: torch.Tensor, tokenizer, remove_punctuation_from_words=False):
	subwords, subword_tokens_list, subword_tokens_indices_list = split_tokens_on_unicode(tokens, tokenizer, remove_punctuation_from_words=remove_punctuation_from_words)
	words = []
	word_tokens = []
	word_tokens_indices = []

	for i, (subword, subword_tokens, subword_tokens_indices) in enumerate(zip(subwords, subword_tokens_list, subword_tokens_indices_list)):
	special = (subword_tokens_indices[0] >= tokenizer.timestamp_begin)
	previous_special = (i > 0) and (subword_tokens_indices_list[i-1][0] >= tokenizer.timestamp_begin)
	next_special = (i < len(subword_tokens_indices_list)-1) and (subword_tokens_indices_list[i+1][0] >= tokenizer.timestamp_begin)
	previous_space = (i > 0) and (not subwords[i-1].strip())
	is_space = not subword.strip()
	with_space = subword.startswith(" ") and not is_space
	punctuation = not is_space and subword.strip() in _punctuation
	if special or (not previous_space and (previous_special or (with_space and not punctuation) or (is_space and not next_special))):
	words.append(subword.strip())
	word_tokens.append(subword_tokens)
	word_tokens_indices.append(subword_tokens_indices)
	else:
	words[-1] = words[-1] + subword.strip()
	word_tokens[-1].extend(subword_tokens)
	word_tokens_indices[-1].extend(subword_tokens_indices)

	return words, word_tokens, word_tokens_indices

	def check_vad_method(method, with_version=False):
	if method in [True, "True", "true"]:
	return check_vad_method("silero") # default method
	elif method in [False, "False", "false"]:
	return False
	elif method.startswith("silero"):
	version = None
	if method != "silero":
	assert method.startswith("silero:"), f"Got unexpected VAD method {method}"
	version = method.split(":")[1]
	if not version.startswith("v"):
	version = "v" + version
	try:
	assert float(version[1:]) >= 1
	except:
	raise ValueError(f"Got unexpected silero version {version} (please check https://github.com/snakers4/silero-vad/wiki/Version-history-and-Available-Models)")
	if with_version:
	return ("silero", version)
	else:
	return method
	elif method == "auditok":
	try:
	import auditok
	except ImportError:
	raise ImportError("Please install auditok to use the auditok VAD (or use another VAD method)")
	else:
	raise ValueError(f"Got unexpected VAD method {method}")
	return method

	_silero_vad_model = None
	_has_onnx = None
	def get_vad_segments(audio,
	output_sample=False,
	min_speech_duration=0.1,
	min_silence_duration=0.1,
	dilatation=0.5,
	method="silero",
	):
	"""
	Get speech segments from audio using Silero VAD
	parameters:
	audio: torch.Tensor
	audio data in 16kHz
	output_sample: bool
	if True, return start and end in samples instead of seconds
	min_speech_duration: float
	minimum duration (in sec) of a speech segment
	min_silence_duration: float
	minimum duration (in sec) of a silence segment
	dilatation: float
	how much (in sec) to enlarge each speech segment detected by the VAD
	method: str
	VAD method to use (auditok, silero, silero:v3.1)
	"""
	global _silero_vad_model, _silero_get_speech_ts, _has_onnx

	if method.startswith("silero"):

	version = None
	_, version = check_vad_method(method, True)
	# See discussion https://github.com/linto-ai/whisper-timestamped/pull/142/files#r1398326287
	need_folder_hack = version and (version < "v4")

	if _silero_vad_model is None:
	# ONNX support since 3.1 in silero
	if (version is None or version >= "v3.1") and (_has_onnx is not False):
	onnx=True
	try:
	import onnxruntime
	onnxruntime.set_default_logger_severity(3) # Remove warning "Removing initializer 'XXX'. It is not used by any node and should be removed from the model."
	_has_onnx = True
	except ImportError as err:
	logger.warning(f"Please install onnxruntime to use more efficiently silero VAD")
	_has_onnx = False
	onnx=False
	else:
	onnx=False

	# Choose silero version because of problems with version 4, see https://github.com/linto-ai/whisper-timestamped/issues/74
	repo_or_dir_master = os.path.expanduser("~/.cache/torch/hub/snakers4_silero-vad_master")
	repo_or_dir_specific = os.path.expanduser(f"~/.cache/torch/hub/snakers4_silero-vad_{version}") if version else repo_or_dir_master
	repo_or_dir = repo_or_dir_specific
	tmp_folder = None
	def apply_folder_hack():
	nonlocal tmp_folder
	if os.path.exists(repo_or_dir_master):
	tmp_folder = repo_or_dir_master + ".tmp"
	shutil.move(repo_or_dir_master, tmp_folder)
	# Make a symlink to the v3.1 model, otherwise it fails
	input_exists = os.path.exists(repo_or_dir_specific)
	if not input_exists:
	# Make dummy file for the symlink to work
	os.makedirs(repo_or_dir_specific, exist_ok=True)
	os.symlink(repo_or_dir_specific, repo_or_dir_master)
	if not input_exists:
	shutil.rmtree(repo_or_dir_specific)

	source = "local"
	if not os.path.exists(repo_or_dir):
	# Load specific version of silero
	repo_or_dir = f"snakers4/silero-vad:{version}" if version else "snakers4/silero-vad"
	source = "github"
	if need_folder_hack:
	apply_folder_hack()
	try:
	_silero_vad_model, utils = torch.hub.load(repo_or_dir=repo_or_dir, model="silero_vad", onnx=onnx, source=source)
	except ImportError as err:
	raise RuntimeError(f"Please install what is needed to use the silero VAD (or use another VAD method)") from err
	except Exception as err:
	raise RuntimeError(f"Problem when installing silero with version {version}. Check versions here: https://github.com/snakers4/silero-vad/wiki/Version-history-and-Available-Models") from err
	finally:
	if need_folder_hack:
	if os.path.exists(repo_or_dir_master):
	os.remove(repo_or_dir_master)
	if tmp_folder:
	shutil.move(tmp_folder, repo_or_dir_master)
	assert os.path.isdir(repo_or_dir_specific), f"Unexpected situation: missing {repo_or_dir_specific}"

	_silero_get_speech_ts = utils[0]

	# Cheap normalization of the volume
	audio = audio / max(0.1, audio.abs().max())

	segments = _silero_get_speech_ts(audio, _silero_vad_model,
	min_speech_duration_ms = round(min_speech_duration * 1000),
	min_silence_duration_ms = round(min_silence_duration * 1000),
	return_seconds = False,
	)

	elif method == "auditok":
	import auditok

	# Cheap normalization of the volume
	audio = audio / max(0.1, audio.abs().max())

	data = (audio.numpy() * 32767).astype(np.int16).tobytes()

	segments = auditok.split(
	data,
	sampling_rate=SAMPLE_RATE, # sampling frequency in Hz
	channels=1, # number of channels
	sample_width=2, # number of bytes per sample
	min_dur=min_speech_duration, # minimum duration of a valid audio event in seconds
	max_dur=len(audio)/SAMPLE_RATE, # maximum duration of an event
	max_silence=min_silence_duration, # maximum duration of tolerated continuous silence within an event
	energy_threshold=50,
	drop_trailing_silence=True,
	)

	segments = [{"start": s._meta.start * SAMPLE_RATE, "end": s._meta.end * SAMPLE_RATE} for s in segments]

	else:
	raise ValueError(f"Got unexpected VAD method {method}")

	if dilatation > 0:
	dilatation = round(dilatation * SAMPLE_RATE)
	new_segments = []
	for seg in segments:
	new_seg = {
	"start": max(0, seg["start"] - dilatation),
	"end": min(len(audio), seg["end"] + dilatation)
	}
	if len(new_segments) > 0 and new_segments[-1]["end"] >= new_seg["start"]:
	new_segments[-1]["end"] = new_seg["end"]
	else:
	new_segments.append(new_seg)
	segments = new_segments

	ratio = 1 if output_sample else 1 / SAMPLE_RATE

	if ratio != 1:
	for seg in segments:
	seg["start"] *= ratio
	seg["end"] *= ratio
	if output_sample:
	for seg in segments:
	seg["start"] = round(seg["start"])
	seg["end"] = round(seg["end"])
	return segments

	def remove_non_speech(audio,
	use_sample=False,
	min_speech_duration=0.1,
	min_silence_duration=1,
	method="silero",
	plot=False,
	):
	"""
	Remove non-speech segments from audio (using Silero VAD),
	glue the speech segments together and return the result along with
	a function to convert timestamps from the new audio to the original audio

	parameters:
	audio: torch.Tensor
	audio data in 16kHz
	use_sample: bool
	if True, return start and end in samples instead of seconds
	min_speech_duration: float
	minimum duration (in sec) of a speech segment
	min_silence_duration: float
	minimum duration (in sec) of a silence segment
	method: str
	method to use to remove non-speech segments
	plot: bool or str
	if True, plot the result.
	If a string, save the plot to the given file
	"""

	segments = get_vad_segments(
	audio,
	output_sample=True,
	min_speech_duration=min_speech_duration,
	min_silence_duration=min_silence_duration,
	method=method,
	)

	segments = [(seg["start"], seg["end"]) for seg in segments]
	if len(segments) == 0:
	segments = [(0, audio.shape[-1])]

	audio_speech = torch.cat([audio[..., s:e] for s,e in segments], dim=-1)

	if plot:
	import matplotlib.pyplot as plt
	plt.figure()
	max_num_samples = 10000
	step = (audio.shape[-1] // max_num_samples) + 1
	times = [i*step/SAMPLE_RATE for i in range((audio.shape[-1]-1) // step + 1)]
	plt.plot(times, audio[::step])
	for s, e in segments:
	plt.axvspan(s/SAMPLE_RATE, e/SAMPLE_RATE, color='red', alpha=0.1)
	if isinstance(plot, str):
	plt.savefig(f"{plot}.VAD.jpg", bbox_inches='tight', pad_inches=0)
	else:
	plt.show()

	if not use_sample:
	segments = [(float(s)/SAMPLE_RATE, float(e)/SAMPLE_RATE) for s,e in segments]

	return audio_speech, lambda t, t2 = None: do_convert_timestamps(segments, t, t2)

	def do_convert_timestamps(segments, t, t2 = None):
	"""
	Convert timestamp from audio without non-speech segments to original audio (with non-speech segments)

	parameters:
	segments: list of tuple (start, end) corresponding to non-speech segments in original audio
	t: timestamp to convert
	t2: second timestamp to convert (optional), when the two timestamps should be in the same segment
	"""
	assert len(segments)
	ioffset = 0 # Input offset
	ooffset = 0 # Output offset
	ipreviousend = 0
	result = []
	for istart, iend in segments:
	ostart = ooffset
	oend = ostart + (iend - istart)
	ooffset = oend
	ioffset += istart - ipreviousend
	ipreviousend = iend
	t_in = t <= oend
	t2_in = t_in if t2 is None else t2 <= oend
	if t_in or t2_in:
	result.append([
	max(istart, min(iend, ioffset + t)),
	max(istart, min(iend, ioffset + t2)) if t2 is not None else None
	])
	if t_in and t2_in:
	break
	if not len(result):
	result.append(
	[ioffset + t, ioffset + t2 if t2 is not None else None]
	)

	if len(result) > 1:
	# Minimize difference between durations
	result = sorted(result, key=lambda x: abs(abs(t2-t) - abs(x[1]-x[0])))
	result = result[0]
	if t2 is None:
	result = round(result[0], 2)
	else:
	result = [round(x, 2) for x in result]
	return result

	def remove_last_null_duration_words(transcription, words, recompute_text=False):
	"""
	Remove words with null duration happening at the end of a chunk (probable Whisper hallucinations)
	"""
	# First group segments by audio chunk
	segments_groups = {}
	seek = None
	current_chunk = -1
	for i, segment in enumerate(transcription["segments"]):
	if segment["seek"] != seek:
	current_chunk += 1
	seek = segment["seek"]
	segments_groups[i] = current_chunk

	# Remove words with null duration happening at the end of a chunk
	current_chunk = -1
	is_last_empty = False
	to_remove = []
	for i, word in enumerate(words[::-1]): # Reverse order
	i = len(words) - i - 1
	empty = (word["start"] == word["end"])
	idx_segment = word["idx_segment"]
	group = segments_groups[idx_segment]
	if current_chunk != group:
	is_last_empty = empty
	current_chunk = group
	elif not empty:
	is_last_empty = False
	if is_last_empty:
	# Remove word
	to_remove.append(i)
	# Shorten text of segment
	full_word = "".join(word["tokens"])
	logger.debug(f"Removing word {i+1}/{len(words)} \"{full_word}\" with empty duration at the end of segment {idx_segment+1}/{len(transcription['segments'])}")
	segment = transcription["segments"][idx_segment]
	text = segment["text"]
	if not text.endswith(full_word): # see issue #62
	if text.endswith(full_word[:-1]):
	full_word = full_word[:-1]
	elif text[:-1].endswith(full_word):
	text = text[:-1]
	else:
	raise RuntimeError(f"\"{text}\" not ending with \"{full_word}\"")
	text = text[:-len(full_word)]
	if i > 0 and words[i-1]["idx_segment"] == idx_segment:
	segment["text"] = text
	else:
	logger.debug(f"Removing empty segment {idx_segment}")
	# Remove segment with no more words
	transcription["segments"].pop(idx_segment)
	for j in range(i+1, len(words)):
	words[j]["idx_segment"] -= 1
	recompute_text = True

	for i in to_remove:
	words.pop(i) # Warning: inplace modification

	if recompute_text:
	transcription["text"] = "".join([s["text"] for s in transcription["segments"]])

	return transcription, words


	def ensure_increasing_positions(segments, min_duration=0):
	"""
	Ensure that "start" and "end" come in increasing order
	"""
	has_modified_backward = False
	previous_end = 0
	for i, seg in enumerate(segments):
	if seg["start"] < previous_end:
	assert i > 0
	new_start = round_timestamp((previous_end + seg["start"]) / 2)
	if new_start < segments[i-1]["start"] + min_duration:
	new_start = previous_end
	else:
	segments[i-1]["end"] = new_start
	has_modified_backward = True
	seg["start"] = new_start
	if seg["end"] <= seg["start"] + min_duration:
	seg["end"] = seg["start"] + min_duration
	previous_end = seg["end"]
	if has_modified_backward:
	return ensure_increasing_positions(segments, min_duration)

	previous_end = 0
	for seg in segments:
	seg["start"] = round_timestamp(seg["start"])
	seg["end"] = round_timestamp(seg["end"])
	assert seg["start"] >= previous_end, f"Got segment {seg} coming before the previous finishes ({previous_end} > {seg['start']})"
	assert seg["end"] >= seg["start"], f"Got segment {seg} with end < start"
	previous_end = seg["end"]

	return segments

	## Some utilities for writing transcripts to files

	def flatten(list_of_lists, key = None):
	for sublist in list_of_lists:
	for item in sublist.get(key, []) if key else sublist:
	yield item

	def remove_keys(list_of_dicts, key):
	for d in list_of_dicts:
	yield {k: d[k] for k in d.keys() - {key}}


	def write_csv(transcript, file, sep = ",", text_first=True, format_timestamps=None, header=False):
	writer = csv.writer(file, delimiter=sep)
	if format_timestamps is None: format_timestamps = lambda x: x
	if header is True:
	header = ["text", "start", "end"] if text_first else ["start", "end", "text"]
	if header:
	writer.writerow(header)
	if text_first:
	writer.writerows(
	[[segment["text"].strip(), format_timestamps(segment["start"]), format_timestamps(segment["end"])] for segment in transcript]
	)
	else:
	writer.writerows(
	[[format_timestamps(segment["start"]), format_timestamps(segment["end"]), segment["text"].strip()] for segment in transcript]
	)

	# https://stackoverflow.com/questions/66588715/runtimeerror-cudnn-error-cudnn-status-not-initialized-using-pytorch
	# CUDA initialization may fail on old GPU card
	def force_cudnn_initialization(device=None, s=32):
	if device is None:
	device = get_default_device()
	torch.nn.functional.conv2d(torch.zeros(s, s, s, s, device=device), torch.zeros(s, s, s, s, device=device))

	def get_default_device():
	if torch.cuda.is_available():
	device = "cuda"
	elif find_spec('torch.xpu') is not None and torch.xpu.is_available():
	device = "xpu"
	else:
	device = "cpu"
	return device

	# base85-encoded (n_layers, n_heads) boolean arrays indicating the cross-attention heads that are
	# highly correlated to the word-level timing, i.e. the alignment between audio and text tokens.
	_ALIGNMENT_HEADS = {
	"tiny.en": b"ABzY8J1N>@0{>%R00Bk>$p{7v037`oCl~+#00",
	"tiny": b"ABzY8bu8Lr0{>%RKn9Fp%m@SkK7Kt=7ytkO",
	"base.en": b"ABzY8;40c<0{>%RzzG;p*o+Vo09\|#PsxSZm00",
	"base": b"ABzY8KQ!870{>%RzyTQH3`Q^yNP!>##QT-<FaQ7m",
	"small.en": b"ABzY8>?_)10{>%RpeA61k&I\|OI3I$65C{;;pbCHh0B{qLQ;+}v00",
	"small": b"ABzY8DmU6=0{>%Rpa?J`kvJ6qF(V^F86#Xh7JUGMK}P<N0000",
	"medium.en": b"ABzY8usPae0{>%R7<zz_OvQ{)4kMa0BMw6u5rT}kRKX;$NfYBv00*Hl@qhsU00",
	"medium": b"ABzY8B0Jh+0{>%R7}kK1fFL7w6%<-Pf*t^=N)Qr&0RR9",
	"large-v1": b"ABzY8r9j$a0{>%R7#4sLmoOs{s)o3~84-RPdcFk!JR<kSfC2yj",
	"large-v2": b'ABzY8zd+h!0{>%R7=D0pU<_bnWW*tkYAhobTNnu$jnkEkXqp)j;w1Tzk)UH3X%SZd&fFZ2fC2yj',
	"large-v3": b"ABzY8gWO1E0{>%R7(9S+Kn!D~%ngiGaR?*L!iJG9p-nab0JQ=-{D1-g00",
	}

	_PARAMETERS_TO_MODEL_NAME = {
	37184256 : "tiny.en",
	37184640 : "tiny",
	71825408 : "base.en",
	71825920 : "base",
	240582144 : "small.en",
	240582912 : "small",
	762320896 : "medium.en",
	762321920 : "medium",
	1541384960 : "large",
	1541570560 : "large-v3",
	}

	def get_alignment_heads(model, max_top_layer=3):
	if hasattr(model, "alignment_heads"): # Since version 20230306
	return model.alignment_heads
	num_parameters = _get_number_of_parameters(model)
	num_layers = model.dims.n_text_layer
	num_heads = model.dims.n_text_head
	if num_parameters not in _PARAMETERS_TO_MODEL_NAME:
	logger.warning("Could not retrieve alignment heads : taking all attention heads from the top layers")
	return None
	model_name = _PARAMETERS_TO_MODEL_NAME[num_parameters]
	if model_name == "large":
	if next(model.parameters())[0,0,0] > 0:
	model_name = "large-v1"
	else:
	model_name = "large-v2"
	return _get_alignment_heads(model_name, num_layers, num_heads)

	def _get_alignment_heads(model_name, num_layers, num_heads):
	dump = _ALIGNMENT_HEADS[model_name]
	array = np.frombuffer(gzip.decompress(base64.b85decode(dump)), dtype=bool).copy()
	mask = torch.from_numpy(array).reshape(num_layers, num_heads)
	alignment_heads = mask.to_sparse()
	return alignment_heads

	def _get_number_of_parameters(model):
	return sum(p.numel() for p in model.parameters())

	from typing import Optional, Union
	def load_model(
	name: str,
	device: Optional[Union[str, torch.device]] = None,
	download_root: str = None,
	in_memory: bool = False,
	):
	extension = os.path.splitext(name)[-1] if os.path.isfile(name) else None

	if name in whisper.available_models() or extension == ".pt":
	return whisper.load_model(name, device=device, download_root=download_root, in_memory=in_memory)

	# Otherwise, assume transformers
	if extension in [".ckpt", ".bin"]:
	model_path = name
	else:
	# Search for the cached file (download if necessary)
	try:
	import transformers
	except ImportError:
	raise ImportError(f"If you are trying to download a HuggingFace model with {name}, please install first the transformers library")
	from transformers.utils import cached_file

	try:
	model_path = cached_file(name, "pytorch_model.bin", cache_dir=download_root, use_auth_token=None, revision=None)
	except Exception as e:
	try:
	if isinstance(e, OSError):
	model_path = cached_file(name, "whisper.ckpt", cache_dir=download_root, use_auth_token=None, revision=None)
	else:
	raise e
	except:
	raise RuntimeError(f"Original error: {e}\nCould not find model {name} from HuggingFace nor local folders.")
	# Load HF Model
	hf_state_dict = torch.load(model_path, map_location="cpu")

	# Rename layers
	for key in list(hf_state_dict.keys())[:]:
	new_key = hf_to_whisper_states(key)
	if new_key is None:
	hf_state_dict.pop(key)
	elif new_key != key:
	hf_state_dict[new_key] = hf_state_dict.pop(key)


	# Init Whisper Model and replace model weights
	dims = whisper.model.ModelDimensions(**states_to_dim(hf_state_dict))

	if "proj_out.weight" in hf_state_dict:
	hf_state_dict["decoder.proj_out.weight"] = hf_state_dict.pop("proj_out.weight")
	logger.warning("Using untied projection layer")
	whisper_model = WhisperUntied(dims)
	else:
	whisper_model = whisper.model.Whisper(dims)

	whisper_model.load_state_dict(hf_state_dict)
	del hf_state_dict
	if hasattr(whisper_model, "alignment_heads"):
	del whisper_model.alignment_heads # Will be recomputed later
	whisper_model = whisper_model.to(device)
	return whisper_model

	# Credit: https://github.com/openai/whisper/discussions/830
	def hf_to_whisper_states(text):
	# From Speechbrain
	if text == "_mel_filters":
	return None

	# From PEFT
	if "default" in text:
	# print(f"WARNING: Ignoring {text}")
	return None
	if text.startswith("base_model.model."):
	text = text[len("base_model.model."):]

	text = re.sub('.layers.', '.blocks.', text)
	text = re.sub('.self_attn.', '.attn.', text)
	text = re.sub('.q_proj.', '.query.', text)
	text = re.sub('.k_proj.', '.key.', text)
	text = re.sub('.v_proj.', '.value.', text)
	text = re.sub('.out_proj.', '.out.', text)
	text = re.sub('.fc1.', '.mlp.0.', text)
	text = re.sub('.fc2.', '.mlp.2.', text)
	text = re.sub('.fc3.', '.mlp.3.', text)
	text = re.sub('.fc3.', '.mlp.3.', text)
	text = re.sub('.encoder_attn.', '.cross_attn.', text)
	text = re.sub('.cross_attn.ln.', '.cross_attn_ln.', text)
	text = re.sub('.embed_positions.weight', '.positional_embedding', text)
	text = re.sub('.embed_tokens.', '.token_embedding.', text)
	text = re.sub('model.', '', text)
	text = re.sub('attn.layer_norm.', 'attn_ln.', text)
	text = re.sub('.final_layer_norm.', '.mlp_ln.', text)
	text = re.sub('encoder.layer_norm.', 'encoder.ln_post.', text)
	text = re.sub('decoder.layer_norm.', 'decoder.ln.', text)
	return text

	def states_to_dim(state_dict):
	n_audio_state = len(state_dict['encoder.ln_post.bias'])
	n_text_state = len(state_dict["decoder.ln.bias"])
	return {
	"n_mels": state_dict["encoder.conv1.weight"].shape[1], # 80
	"n_vocab": state_dict["decoder.token_embedding.weight"].shape[0], # 51864 / 51865
	"n_audio_ctx": state_dict["encoder.positional_embedding"].shape[0], # 1500
	"n_audio_state": n_audio_state, # 384 / 512 / 768 / 1024 / 1280
	"n_audio_head": n_audio_state // 64, # 6 / 8 / 12 / 16 / 20
	"n_audio_layer": len(set([".".join(k.split(".")[:3]) for k in state_dict.keys() if "encoder.blocks." in k])), # 4 / 6 / 12 / 24 / 32
	"n_text_ctx": state_dict["decoder.positional_embedding"].shape[0], # 448
	"n_text_state": n_text_state, # 384 / 512 / 768 / 1024 / 1280
	"n_text_head": n_text_state // 64, # 6 / 8 / 12 / 16 / 20
	"n_text_layer": len(set([".".join(k.split(".")[:3]) for k in state_dict.keys() if "decoder.blocks." in k])), # 4 / 6 / 12 / 24 / 32
	}

	class TextDecoderUntied(whisper.model.TextDecoder):
	"""
	Same as TextDecoder but with untied weights
	"""
	def __init__(self, args, *kwargs):
	import torch
	super().__init__(args, *kwargs)

	n_vocab, n_state = self.token_embedding.weight.shape

	self.proj_out = torch.nn.Linear(n_state, n_vocab, bias=False)

	def forward(self, x, xa, kv_cache = None):
	offset = next(iter(kv_cache.values())).shape[1] if kv_cache else 0
	x = self.token_embedding(x) + self.positional_embedding[offset : offset + x.shape[-1]]
	x = x.to(xa.dtype)

	for block in self.blocks:
	x = block(x, xa, mask=self.mask, kv_cache=kv_cache)

	x = self.ln(x)

	# logits = self.proj_out(x).float()
	# logits = (x @ torch.transpose(self.proj_out.weight.to(x.dtype), 0, 1)).float()
	logits = self.proj_out.to(x.dtype)(x).float()

	return logits

	class WhisperUntied(whisper.model.Whisper):
	def __init__(self, args, *kwargs):
	super().__init__(args, *kwargs)
	self.decoder = TextDecoderUntied(
	self.dims.n_vocab,
	self.dims.n_text_ctx,
	self.dims.n_text_state,
	self.dims.n_text_head,
	self.dims.n_text_layer,
	)

	def cli():

	import os
	import sys
	import argparse
	import json

	from whisper.utils import str2bool, optional_float, optional_int

	try:
	# Old whisper version # Before https://github.com/openai/whisper/commit/da600abd2b296a5450770b872c3765d0a5a5c769
	from whisper.utils import write_txt, write_srt, write_vtt
	write_tsv = lambda transcript, file: write_csv(transcript, file, sep="\t", header=True, text_first=False, format_timestamps=lambda x: round(1000 * x))

	except ImportError:
	# New whisper version
	from whisper.utils import get_writer

	def do_write(transcript, file, output_format):
	writer = get_writer(output_format, os.path.curdir)
	try:
	return writer.write_result({"segments": transcript}, file)
	except TypeError:
	# Version > 20230314
	return writer.write_result({"segments": list(transcript)}, file, {
	"highlight_words": False,
	"max_line_width": None,
	"max_line_count": None,
	})
	def get_do_write(output_format):
	return lambda transcript, file: do_write(transcript, file, output_format)

	write_txt = get_do_write("txt")
	write_srt = get_do_write("srt")
	write_vtt = get_do_write("vtt")
	write_tsv = get_do_write("tsv")

	parser = argparse.ArgumentParser(
	description='Transcribe a single audio with whisper and compute word timestamps',
	formatter_class=argparse.ArgumentDefaultsHelpFormatter
	)
	parser.add_argument('-v', '--version', help="show version and exit", action='version', version=f'{__version__}')
	parser.add_argument('--versions', help="show versions (of whisper-timestamped and whisper) and exit", action='version',
	version=f'{__version__} -- Whisper {whisper.__version__} in {os.path.realpath(os.path.dirname(whisper.__file__))}')

	parser.add_argument('audio', help="audio file(s) to transcribe", nargs='+')
	parser.add_argument('--model', help=f"name of the Whisper model to use. Examples: {', '.join(whisper.available_models())}", default="small")
	parser.add_argument("--model_dir", default=None, help="the path to save model files; uses ~/.cache/whisper by default", type=str)
	parser.add_argument("--device", default=get_default_device(), help="device to use for PyTorch inference")
	parser.add_argument("--output_dir", "-o", default=None, help="directory to save the outputs", type=str)
	valid_formats = ["txt", "vtt", "srt", "tsv", "csv", "json"]
	def str2output_formats(string):
	if string == "all":
	return valid_formats
	formats = string.split(",")
	for format in formats:
	if format not in valid_formats:
	raise ValueError(f"Expected one of {valid_formats}, got {format}")
	return formats
	parser.add_argument("--output_format", "-f", default="all", help=f"Format(s) of the output file(s). Possible formats are: {', '.join(valid_formats)}. Several formats can be specified by using commas (ex: \"json,vtt,srt\"). By default (\"all\"), all available formats will be produced", type=str2output_formats)

	parser.add_argument("--task", default="transcribe", help="whether to perform X->X speech recognition ('transcribe') or X->English translation ('translate')", choices=["transcribe", "translate"], type=str)
	parser.add_argument('--language', help=f"language spoken in the audio, specify None to perform language detection.", choices=sorted(whisper.tokenizer.LANGUAGES.keys()) + sorted([k.title() for k in whisper.tokenizer.TO_LANGUAGE_CODE.keys()]), default=None)
	# f"{', '.join(sorted(k+'('+v+')' for k,v in whisper.tokenizer.LANGUAGES.items()))}

	parser.add_argument('--vad', default=False, help="whether to run Voice Activity Detection (VAD) to remove non-speech segment before applying Whisper model (removes hallucinations). Can be: True, False, silero, silero:3.1 (or another version), or autitok. Some additional libraries might be needed")
	parser.add_argument('--detect_disfluencies', default=False, help="whether to try to detect disfluencies, marking them as special words [*]", type=str2bool)
	parser.add_argument('--recompute_all_timestamps', default=not TRUST_WHISPER_TIMESTAMP_BY_DEFAULT, help="Do not rely at all on Whisper timestamps (Experimental option: did not bring any improvement, but could be useful in cases where Whipser segment timestamp are wrong by more than 0.5 seconds)", type=str2bool)
	parser.add_argument("--punctuations_with_words", default=True, help="whether to include punctuations in the words", type=str2bool)

	parser.add_argument("--temperature", default=0.0, help="temperature to use for sampling", type=float)
	parser.add_argument("--best_of", type=optional_int, default=None if USE_EFFICIENT_BY_DEFAULT else 5, help="number of candidates when sampling with non-zero temperature")
	parser.add_argument("--beam_size", type=optional_int, default=None if USE_EFFICIENT_BY_DEFAULT else 5, help="number of beams in beam search, only applicable when temperature is zero")
	parser.add_argument("--patience", type=float, default=None, help="optional patience value to use in beam decoding, as in https://arxiv.org/abs/2204.05424, the default (1.0) is equivalent to conventional beam search")
	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 length normalization by default")

	parser.add_argument("--suppress_tokens", default="-1", help="comma-separated list of token ids to suppress during sampling; '-1' will suppress most special characters except common punctuations", type=str)
	parser.add_argument("--initial_prompt", default=None, help="optional text to provide as a prompt for the first window.", type=str)
	parser.add_argument("--condition_on_previous_text", 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", type=str2bool)
	parser.add_argument("--fp16", default=None, help="whether to perform inference in fp16; Automatic by default (True if GPU available, False otherwise)", type=str2bool)

	parser.add_argument("--temperature_increment_on_fallback", default=0.0 if USE_EFFICIENT_BY_DEFAULT else 0.2, help="temperature to increase when falling back when the decoding fails to meet either of the thresholds below", type=optional_float)
	parser.add_argument("--compression_ratio_threshold", default=2.4, help="if the gzip compression ratio is higher than this value, treat the decoding as failed", type=optional_float)
	parser.add_argument("--logprob_threshold", default=-1.0, help="if the average log probability is lower than this value, treat the decoding as failed", type=optional_float)
	parser.add_argument("--no_speech_threshold", default=0.6, help="if the probability of the <\|nospeech\|> token is higher than this value AND the decoding has failed due to `logprob_threshold`, consider the segment as silence", type=optional_float)
	parser.add_argument("--threads", default=0, help="number of threads used by torch for CPU inference; supercedes MKL_NUM_THREADS/OMP_NUM_THREADS", type=optional_int)

	parser.add_argument("--compute_confidence", default=True, help="whether to compute confidence scores for words", type=str2bool)
	parser.add_argument("--verbose", type=str2bool, default=False, help="whether to print out the progress and debug messages of Whisper")
	parser.add_argument('--plot', help="plot word alignments (save the figures if an --output_dir is specified, otherwhise just show figures that have to be closed to continue)", default=False, action="store_true")
	parser.add_argument('--debug', help="print some debug information about word alignement", default=False, action="store_true")

	class ActionSetAccurate(argparse.Action):
	def __init__(self, option_strings, dest, nargs=None, **kwargs):
	assert nargs is None
	super().__init__(option_strings, dest, nargs=0, **kwargs)
	def __call__(self, parser, namespace, values, option_string=None):
	setattr(namespace, "best_of", 5)
	setattr(namespace, "beam_size", 5)
	setattr(namespace, "temperature_increment_on_fallback", 0.2)
	parser.add_argument('--accurate', help="Shortcut to use the same default option as in Whisper (best_of=5, beam_search=5, temperature_increment_on_fallback=0.2)", action=ActionSetAccurate)

	class ActionSetEfficient(argparse.Action):
	def __init__(self, option_strings, dest, nargs=None, **kwargs):
	assert nargs is None
	super().__init__(option_strings, dest, nargs=0, **kwargs)
	def __call__(self, parser, namespace, values, option_string=None):
	setattr(namespace, "best_of", None)
	setattr(namespace, "beam_size", None)
	setattr(namespace, "temperature_increment_on_fallback", None)
	parser.add_argument('--efficient', help="Shortcut to disable beam size and options that requires to sample several times, for an efficient decoding", action=ActionSetEfficient)

	parser.add_argument('--naive', help="use naive approach, doing inference twice (once to get the transcription, once to get word timestamps and confidence scores).", default=False, action="store_true")

	args = parser.parse_args().__dict__
	args.pop("accurate")
	args.pop("efficient")

	temperature = args.pop("temperature")
	temperature_increment_on_fallback = args.pop("temperature_increment_on_fallback")
	if temperature_increment_on_fallback:
	temperature = tuple(np.arange(temperature, 1.0 + 1e-6, temperature_increment_on_fallback))
	else:
	temperature = [temperature]

	threads = args.pop("threads")
	if threads:
	torch.set_num_threads(threads)

	audio_files = args.pop("audio")

	model = args.pop("model")
	device = args.pop("device")
	model_dir = args.pop("model_dir")

	if device.lower().startswith("cuda"):
	force_cudnn_initialization(device)

	output_format = args.pop("output_format")

	model = load_model(model, device=device, download_root=model_dir)

	plot_word_alignment = args.pop("plot")

	debug = args.pop("debug")
	logging.basicConfig()
	if debug:
	logger.setLevel(logging.DEBUG)
	# This supposes to plug a logger with name "WHISPER" into Whisper source code (no harm if it's not set)
	logging.getLogger("WHISPER").setLevel(logging.DEBUG)

	output_dir = args.pop("output_dir")
	if output_dir and not os.path.isdir(output_dir):
	os.makedirs(output_dir)

	args["naive_approach"] = args.pop("naive")
	args["remove_punctuation_from_words"] = not args.pop("punctuations_with_words")
	args["compute_word_confidence"] = args.pop("compute_confidence")
	args["trust_whisper_timestamps"] = not args.pop("recompute_all_timestamps")

	for audio_path in audio_files:

	outname = os.path.join(output_dir, os.path.basename(audio_path)) if output_dir else None

	result = transcribe_timestamped(
	model, audio_path,
	temperature=temperature,
	plot_word_alignment=outname if (outname and plot_word_alignment) else plot_word_alignment,
	**args
	)

	if output_dir:

	if "json" in output_format:
	# save JSON
	with open(outname + ".words.json", "w", encoding="utf-8") as js:
	json.dump(result, js, indent=2, ensure_ascii=False)

	# save TXT
	if "txt" in output_format:
	with open(outname + ".txt", "w", encoding="utf-8") as txt:
	write_txt(result["segments"], file=txt)

	# save VTT
	if "vtt" in output_format:
	with open(outname + ".vtt", "w", encoding="utf-8") as vtt:
	write_vtt(remove_keys(result["segments"], "words"), file=vtt)
	with open(outname + ".words.vtt", "w", encoding="utf-8") as vtt:
	write_vtt(flatten(result["segments"], "words"), file=vtt)

	# save SRT
	if "srt" in output_format:
	with open(outname + ".srt", "w", encoding="utf-8") as srt:
	write_srt(remove_keys(result["segments"], "words"), file=srt)
	with open(outname + ".words.srt", "w", encoding="utf-8") as srt:
	write_srt(flatten(result["segments"], "words"), file=srt)

	# save CSV
	if "csv" in output_format:
	with open(outname + ".csv", "w", encoding="utf-8") as csv:
	write_csv(result["segments"], file=csv)
	with open(outname + ".words.csv", "w", encoding="utf-8") as csv:
	write_csv(flatten(result["segments"], "words"), file=csv)

	# save TSV
	if "tsv" in output_format:
	with open(outname + ".tsv", "w", encoding="utf-8") as csv:
	write_tsv(result["segments"], file=csv)
	with open(outname + ".words.tsv", "w", encoding="utf-8") as csv:
	write_tsv(flatten(result["segments"], "words"), file=csv)

	elif not args["verbose"]:

	json.dump(filtered_keys(result), sys.stdout, indent=2, ensure_ascii=False)


	def filtered_keys(result, keys = [
	"text",
	"segments", "words",
	"language",
	"start",
	"end",
	"confidence",
	"language_probs",
	]):
	if isinstance(result, dict):
	return {k: (filtered_keys(v, keys) if k not in ["language_probs"] else v) for k, v in result.items() if k in keys}
	if isinstance(result, list):
	return [filtered_keys(v, keys) for v in result]
	if isinstance(result, float):
	return round(result, 2)
	return result


	if __name__ == "__main__":
	cli()