@huggingface/transformers
Version:
State-of-the-art Machine Learning for the web. Run 🤗 Transformers directly in your browser, with no need for a server!
515 lines (450 loc) • 22.8 kB
JavaScript
import { cat, mean, Tensor, stack, std_mean } from '../../utils/tensor.js';
import { PreTrainedModel } from '../modeling_utils.js';
import { WhisperGenerationConfig } from './generation_whisper.js';
import { whisper_language_to_code } from './common_whisper.js';
import { prepareTensorForDecode } from '../../tokenization_utils.js';
import {
LogitsProcessorList,
SuppressTokensAtBeginLogitsProcessor,
WhisperTimeStampLogitsProcessor,
} from '../../generation/logits_process.js';
import { medianFilter, dynamic_time_warping } from '../../utils/maths.js';
import { mergeArrays } from '../../utils/core.js';
import { ModelOutput } from '../modeling_outputs.js';
import { logger } from '../../utils/logger.js';
export class WhisperPreTrainedModel extends PreTrainedModel {
requires_attention_mask = false;
main_input_name = 'input_features';
forward_params = [
'input_features',
'attention_mask',
'decoder_input_ids',
'decoder_attention_mask',
'past_key_values',
];
}
/**
* WhisperModel class for training Whisper models without a language model head.
*/
export class WhisperModel extends WhisperPreTrainedModel {}
/**
* WhisperForConditionalGeneration class for generating conditional outputs from Whisper models.
*/
export class WhisperForConditionalGeneration extends WhisperPreTrainedModel {
_prepare_generation_config(generation_config, kwargs) {
return /** @type {WhisperGenerationConfig} */ (
super._prepare_generation_config(generation_config, kwargs, WhisperGenerationConfig)
);
}
/**
*
* @param {WhisperGenerationConfig} generation_config
*/
_retrieve_init_tokens(generation_config) {
// prefix tokens are of the form:
// - Multilingual: <|startoftranscript|> <|lang_id|> <|task|> [<|notimestamps|>]
// - English-only: <|startoftranscript|> [<|notimestamps|>]
// 1. Handle <|startoftranscript|> token
const init_tokens = [generation_config.decoder_start_token_id];
// 2. Handle <|lang_id|> and <|task> tokens
let language = generation_config.language;
const task = generation_config.task;
if (generation_config.is_multilingual) {
if (!language) {
// TODO: Implement language detection
logger.warn('No language specified - defaulting to English (en).');
language = 'en';
}
// Add language token
const language_code = whisper_language_to_code(language);
const language_token = `<|${language_code}|>`;
init_tokens.push(generation_config.lang_to_id[language_token]);
// Add task token
// NOTE: Defaults to 'transcribe' if no task is specified
init_tokens.push(generation_config.task_to_id[task ?? 'transcribe']);
} else if (language || task) {
throw new Error(
'Cannot specify `task` or `language` for an English-only model. If the model is intended to be multilingual, pass `is_multilingual=true` to generate, or update the generation config.',
);
}
// 3. Handle <|notimestamps|> token
if (
!generation_config.return_timestamps &&
generation_config.no_timestamps_token_id &&
init_tokens.at(-1) !== generation_config.no_timestamps_token_id
) {
init_tokens.push(generation_config.no_timestamps_token_id);
} else if (
generation_config.return_timestamps &&
init_tokens.at(-1) === generation_config.no_timestamps_token_id
) {
logger.warn(
'<|notimestamps|> prompt token is removed from generation_config since `return_timestamps` is set to `true`.',
);
init_tokens.pop();
}
// let's make sure we don't pass `null` tokens as prompt tokens
return init_tokens.filter((token) => token != null);
}
/**
* Transcribes or translates log-mel input features to a sequence of auto-regressively generated token ids.
* @param {import('./generation_whisper.js').WhisperGenerationFunctionParameters} options
* @returns {Promise<ModelOutput|Tensor>} The output of the model, which can contain the generated token ids, attentions, and scores.
*/
async generate({
inputs = null,
generation_config = null,
logits_processor = null,
stopping_criteria = null,
// Whisper-specific options (passed to kwargs)
// prompt_ids = null,
// language = null,
// task = null,
...kwargs
}) {
generation_config = this._prepare_generation_config(generation_config, kwargs);
const init_tokens =
kwargs.decoder_input_ids instanceof Tensor
? prepareTensorForDecode(kwargs.decoder_input_ids)
: (kwargs.decoder_input_ids ?? this._retrieve_init_tokens(generation_config));
if (generation_config.return_timestamps) {
logits_processor ??= new LogitsProcessorList();
logits_processor.push(new WhisperTimeStampLogitsProcessor(generation_config, init_tokens));
}
if (generation_config.begin_suppress_tokens) {
logits_processor ??= new LogitsProcessorList();
logits_processor.push(
new SuppressTokensAtBeginLogitsProcessor(generation_config.begin_suppress_tokens, init_tokens.length),
);
}
if (generation_config.return_token_timestamps) {
if (!generation_config.alignment_heads) {
throw new Error(
'Model generation config has no `alignment_heads`, token-level timestamps not available. ' +
'See https://gist.github.com/hollance/42e32852f24243b748ae6bc1f985b13a on how to add this property to the generation config.',
);
}
if (generation_config.task === 'translate') {
logger.warn("Token-level timestamps may not be reliable for task 'translate'.");
}
generation_config.output_attentions = true;
generation_config.return_dict_in_generate = true;
}
// For timestamp mode, use seek-based sequential generation.
// This matches Python's WhisperForConditionalGeneration.generate() which uses a seek loop
// to handle audio that the model doesn't fully transcribe in a single pass.
// Skip the seek loop when max_new_tokens is explicitly set (e.g., for prefix token tests).
if (generation_config.return_timestamps && !kwargs.max_new_tokens) {
return this._generate_with_seek({
inputs,
generation_config,
logits_processor,
init_tokens,
kwargs,
});
}
const outputs = await super.generate({
inputs,
generation_config,
logits_processor,
decoder_input_ids: init_tokens,
...kwargs,
});
if (generation_config.return_token_timestamps) {
outputs['token_timestamps'] = this._extract_token_timestamps(
// @ts-expect-error TS2345
outputs,
generation_config.alignment_heads,
generation_config.num_frames,
0.02,
init_tokens.length,
);
}
return outputs;
}
/**
* Generates with a seek loop for timestamp mode, re-encoding and generating
* for each segment until all audio frames are consumed.
* This matches Python's WhisperForConditionalGeneration.generate() behavior.
* @private
*/
async _generate_with_seek({ inputs, generation_config, logits_processor, init_tokens, kwargs }) {
const timestamp_begin = generation_config.no_timestamps_token_id + 1;
const eos_token_id = Array.isArray(generation_config.eos_token_id)
? generation_config.eos_token_id[0]
: generation_config.eos_token_id;
const return_token_timestamps = generation_config.return_token_timestamps;
// input_features shape: [batch=1, n_mels, total_frames]
const input_features = inputs;
const total_frames = input_features.dims[2];
// The encoder downsamples by input_stride (=2 for whisper), so:
// num_segment_frames = input_stride * max_source_positions = 3000 mel frames per segment
// Timestamp token T maps to mel frame position T * input_stride
const input_stride = 2;
// @ts-expect-error ts(2339)
const max_source_positions = /** @type {number} */ (this.config.max_source_positions);
const num_segment_frames = input_stride * max_source_positions;
let seek = 0;
const allTokens = [];
const allTokenTimestamps = [];
while (seek < total_frames) {
// Slice input features for this segment
const seek_end = Math.min(seek + num_segment_frames, total_frames);
const segment_input = input_features.slice(null, null, [seek, seek_end]);
// Pad to full segment size if needed (whisper expects fixed-length input)
let segment_features;
const segment_frames = segment_input.dims[2];
if (segment_frames < num_segment_frames) {
const n_mels = input_features.dims[1];
const padded_data = new Float32Array(n_mels * num_segment_frames);
const src = /** @type {Float32Array} */ (segment_input.data);
// Copy each mel band row separately to handle the stride difference
for (let m = 0; m < n_mels; ++m) {
padded_data.set(src.subarray(m * segment_frames, (m + 1) * segment_frames), m * num_segment_frames);
}
segment_features = new Tensor('float32', padded_data, [1, n_mels, num_segment_frames]);
} else {
segment_features = segment_input;
}
// Reset logits processor begin_index for each segment
if (logits_processor) {
for (const proc of logits_processor) {
if ('begin_index' in proc) {
proc.begin_index = init_tokens.length;
}
}
}
const outputs = /** @type {any} */ (
await super.generate({
inputs: segment_features,
generation_config,
logits_processor,
decoder_input_ids: init_tokens,
...kwargs,
})
);
// Extract tokens (skip init_tokens prefix)
const raw_sequence = return_token_timestamps ? outputs.sequences : /** @type {Tensor} */ (outputs);
const generated_tokens = raw_sequence[0].tolist().map(Number).slice(init_tokens.length);
// Extract token-level timestamps for this seek pass if needed
let seek_token_timestamps;
if (return_token_timestamps) {
outputs['token_timestamps'] = this._extract_token_timestamps(
outputs,
generation_config.alignment_heads,
Math.floor((seek_end - seek) / input_stride),
0.02,
init_tokens.length,
);
const time_offset = (seek / input_stride) * 0.02;
seek_token_timestamps = outputs.token_timestamps[0]
.tolist()
.slice(init_tokens.length)
.map((/** @type {number} */ t) => t + time_offset);
}
// Remove trailing EOS
if (generated_tokens.length > 0 && generated_tokens.at(-1) === eos_token_id) {
generated_tokens.pop();
}
if (generated_tokens.length === 0) {
// No tokens generated — skip the rest of the audio
break;
}
// Determine seek advancement using the same logic as Python's _retrieve_segment:
// 1. Find consecutive timestamp token pairs (segment boundaries)
// 2. If the sequence ends with a single timestamp (no speech after it),
// consume all remaining frames in this segment
// 3. Otherwise, seek to the last complete segment boundary
const is_timestamp = generated_tokens.map((t) => t >= timestamp_begin);
// Check for single_timestamp_ending: last token is timestamp, second-to-last is not
const single_timestamp_ending =
generated_tokens.length >= 2 &&
is_timestamp[generated_tokens.length - 1] &&
!is_timestamp[generated_tokens.length - 2];
// Find consecutive timestamp pairs (segment boundaries)
const segment_boundary_indices = [];
for (let i = 0; i < generated_tokens.length - 1; ++i) {
if (is_timestamp[i] && is_timestamp[i + 1]) {
segment_boundary_indices.push(i + 1); // index of the second token in the pair
}
}
let segment_offset;
let tokens_to_keep = generated_tokens.length;
if (segment_boundary_indices.length > 0) {
if (single_timestamp_ending) {
// Ends with a single timestamp after the last pair — no more speech
segment_offset = seek_end - seek;
} else {
// Ends mid-segment — seek to the last pair's end timestamp
// Discard tokens after the last pair (they're from an incomplete segment)
// Keep up to the first token of the last pair (the end-of-segment timestamp),
// excluding the second token (the start-of-next-segment marker)
const last_boundary = segment_boundary_indices.at(-1);
const last_ts_pos = generated_tokens[last_boundary - 1] - timestamp_begin;
segment_offset = last_ts_pos * input_stride;
tokens_to_keep = last_boundary;
}
} else {
// No consecutive pairs found — consume entire segment
segment_offset = seek_end - seek;
}
// Offset timestamp tokens by the current seek position so they're
// monotonically increasing across segments. Cap at the maximum valid
// timestamp token (30.00s = 1500 positions) to stay within the token vocab.
const timestamp_offset = Math.floor(seek / input_stride);
const max_timestamp_token = timestamp_begin + 1500;
for (let i = 0; i < tokens_to_keep; ++i) {
if (generated_tokens[i] >= timestamp_begin) {
generated_tokens[i] = Math.min(generated_tokens[i] + timestamp_offset, max_timestamp_token);
}
}
allTokens.push(...generated_tokens.slice(0, tokens_to_keep));
if (seek_token_timestamps) {
allTokenTimestamps.push(...seek_token_timestamps.slice(0, tokens_to_keep));
}
seek += segment_offset;
}
// Add EOS back
allTokens.push(eos_token_id);
// Reconstruct output
const full_sequence = [...init_tokens, ...allTokens];
if (return_token_timestamps) {
// Return dict format with sequences and token_timestamps
const sequences = new Tensor('int64', full_sequence.map(BigInt), [1, full_sequence.length]);
// Pad token_timestamps to match full_sequence (init_tokens get 0.0)
const full_timestamps = [...new Array(init_tokens.length).fill(0), ...allTokenTimestamps, 0];
const token_timestamps = new Tensor('float32', new Float32Array(full_timestamps), [
1,
full_timestamps.length,
]);
return { sequences, token_timestamps };
}
return new Tensor('int64', full_sequence.map(BigInt), [1, full_sequence.length]);
}
/**
* Calculates token-level timestamps using the encoder-decoder cross-attentions and
* dynamic time-warping (DTW) to map each output token to a position in the input audio.
* If `num_frames` is specified, the encoder-decoder cross-attentions will be cropped before applying DTW.
* @param {Object} generate_outputs Outputs generated by the model
* @param {Tensor[][]} generate_outputs.cross_attentions The cross attentions output by the model
* @param {Tensor} generate_outputs.sequences The sequences output by the model
* @param {number[][]} alignment_heads Alignment heads of the model
* @param {number} [num_frames=null] Number of frames in the input audio.
* @param {number} [time_precision=0.02] Precision of the timestamps in seconds
* @param {number} [num_input_ids=0] Number of decoder input ids (prefix tokens) to skip in DTW
* @returns {Tensor} tensor containing the timestamps in seconds for each predicted token
*/
_extract_token_timestamps(
generate_outputs,
alignment_heads,
num_frames = null,
time_precision = 0.02,
num_input_ids = 0,
) {
if (!generate_outputs.cross_attentions) {
throw new Error(
'Model outputs must contain cross attentions to extract timestamps. ' +
'This is most likely because the model was not exported with `output_attentions=True`.',
);
}
if (num_frames == null) {
logger.warn(
'`num_frames` has not been set, meaning the entire audio will be analyzed. ' +
'This may lead to inaccurate token-level timestamps for short audios (< 30 seconds).',
);
}
// @ts-expect-error TS2339
let median_filter_width = this.config.median_filter_width;
if (median_filter_width === undefined) {
logger.warn('Model config has no `median_filter_width`, using default value of 7.');
median_filter_width = 7;
}
// TODO: Improve batch processing
const batch = generate_outputs.cross_attentions;
// Create a list with `decoder_layers` elements, each a tensor of shape
// (batch size, attention_heads, output length, input length).
const cross_attentions = Array.from(
// @ts-expect-error TS2339
{ length: this.config.decoder_layers },
// Concatenate the cross attentions for each layer across sequence length dimension.
(_, i) =>
cat(
batch.map((x) => x[i]),
2,
),
);
const weights = stack(
alignment_heads.map(([l, h]) => {
if (l >= cross_attentions.length) {
throw new Error(
`Layer index ${l} is out of bounds for cross attentions (length ${cross_attentions.length}).`,
);
}
return num_frames
? cross_attentions[l].slice(null, h, null, [0, num_frames])
: cross_attentions[l].slice(null, h);
}),
).transpose(1, 0, 2, 3);
const [std, calculatedMean] = std_mean(weights, -2, 0, true);
// Normalize and smoothen the weights.
const smoothedWeights = weights.clone(); // [1, 8, seqLength, 1500]
for (let a = 0; a < smoothedWeights.dims[0]; ++a) {
const aTensor = smoothedWeights[a]; // [8, seqLength, 1500]
for (let b = 0; b < aTensor.dims[0]; ++b) {
const bTensor = aTensor[b]; // [seqLength, 1500]
const stdTensorData = std[a][b][0].data; // [1500]
const meanTensorData = calculatedMean[a][b][0].data; // [1500]
for (let c = 0; c < bTensor.dims[0]; ++c) {
let cTensorData = bTensor[c].data; // [1500]
for (let d = 0; d < cTensorData.length; ++d) {
cTensorData[d] = (cTensorData[d] - meanTensorData[d]) / stdTensorData[d];
}
// Apply median filter.
cTensorData.set(medianFilter(cTensorData, median_filter_width));
}
}
}
// Skip decoder_input_ids in the cross-attention weights
const croppedWeights =
num_input_ids > 0
? smoothedWeights.slice(null, null, [num_input_ids, smoothedWeights.dims[2]], null)
: smoothedWeights;
// Average the different cross-attention heads.
const batchedMatrices = [mean(croppedWeights, 1)];
const timestampsShape = generate_outputs.sequences.dims;
const timestamps = new Tensor(
'float32',
new Float32Array(timestampsShape[0] * timestampsShape[1]),
timestampsShape,
);
// Perform dynamic time warping on each element of the batch.
for (let batch_idx = 0; batch_idx < timestampsShape[0]; ++batch_idx) {
// NOTE: Since we run only one batch at a time, we can squeeze to get the same dimensions
// as the python implementation
const matrix = batchedMatrices[batch_idx].neg().squeeze_(0);
const [text_indices, time_indices] = dynamic_time_warping(matrix.tolist());
const diffs = Array.from(
{ length: text_indices.length - 1 },
(v, i) => text_indices[i + 1] - text_indices[i],
);
const jumps = mergeArrays([1], diffs).map((x) => !!x); // convert to boolean
const jump_times = [];
for (let i = 0; i < jumps.length; ++i) {
if (jumps[i]) {
// NOTE: No point in rounding here, since we set to Float32Array later
jump_times.push(time_indices[i] * time_precision);
}
}
// Pad with num_input_ids zeros at the start (for prefix tokens),
// then DTW jump_times, then duplicate last value (for eos token)
const padded = new Array(num_input_ids).fill(0);
padded.push(...jump_times);
if (jump_times.length > 0) {
padded.push(jump_times.at(-1));
}
timestamps[batch_idx].data.set(padded);
}
return timestamps;
}
}
export class LiteWhisperForConditionalGeneration extends WhisperForConditionalGeneration {}