mstar.model.whisper.components.alignment

mstar.model.whisper.components.alignment#

Word-level timestamps from the decoder’s cross-attention.

Whisper was not trained to emit word boundaries; openai-whisper derives them from a handful of cross-attention heads (alignment_heads in the generation config) that track the audio position of the token being predicted. The recipe, reproduced here on M*’s own components:

  1. run the decoder once more over <|sot|><|lang|><|task|><|notimestamps|> text... <|eot|> with teacher forcing and keep those heads’ attention probabilities over the encoder frames that hold audio;

  2. normalize each head over its tokens, smooth along time with a median filter, average the heads: one (tokens, frames) alignment matrix (row i is the step that predicts text token i; the last row predicts end-of-text);

  3. dynamic time warping through that matrix gives a monotonic path; the frame at which the path first enters a row is when that token starts;

  4. tokens are grouped into words (a new word at a leading space or a lone punctuation mark, or at every run of tokens that decodes to whole characters in a language written without spaces); a word starts with its first token and ends where the next word starts, the last one where end-of-text is predicted.

Steps 1 and 2 are GPU work: the decoder submodule’s align walk emits the matrix (alignment_output). Steps 3 and 4 run on the host in the API server, when WhisperModel.postprocess renders that output (word_timings). The functions here are pure tensor/numpy code with no engine dependencies, so the same math is unit-tested on the CPU.

Functions

alignment_matrix(weights, num_frames[, ...])

weights: (heads, rows, frames) attention probabilities of the alignment heads -> the (rows, audio_frames) matrix DTW walks.

alignment_output(weights, tokens, num_frames)

The align walk's one output tensor, from the alignment heads' weights over the whole teacher-forced tokens (<|sot|><|lang|><|task|><|notimestamps|> text <|eot|>).

dtw(cost)

Cheapest monotonic path from the top-left to the bottom-right of cost (rows x frames), moving right, down or diagonally.

median_filter(x[, width])

Sliding median over the last dim (odd width), edges reflected.

merge_punctuation(words)

openai-whisper's merge_punctuations: opening marks join the word after them (scanning right to left), then closing marks join the word before them.

split_words(decode, token_ids[, language])

Group text tokens into words, as openai-whisper's split_to_word_tokens: first into pieces that decode to whole characters (a token ending inside a multi-byte character continues into the next), then, unless language is written without spaces, into words that begin at a leading space or a lone punctuation mark.

token_start_frames(matrix)

The frame at which the DTW path first enters each row of matrix (a high value is a good alignment, so the path minimizes -matrix).

unpack_output(output)

alignment_output's tensor -> (matrix, text token ids, language token id).

warm_up()

Compile the DTW now, in the process that runs it (the API server's data worker, at model load), rather than inside the first word-timestamp request.

word_timings(matrix, text_tokens, decode[, ...])

Words with start / end seconds for one window.
mstar.model.whisper.components.alignment.alignment_matrix(weights, num_frames, medfilt_width=MEDIAN_FILTER_WIDTH)[source]#

weights: (heads, rows, frames) attention probabilities of the alignment heads -> the (rows, audio_frames) matrix DTW walks.

Only the frames that hold audio count (num_frames mel frames, two per encoder position); each head is standardized over its rows so heads with different sharpness weigh alike, smoothed along time, then averaged.

Parameters:
Return type:

Tensor

mstar.model.whisper.components.alignment.alignment_output(weights, tokens, num_frames)[source]#

The align walk’s one output tensor, from the alignment heads’ weights over the whole teacher-forced tokens (<|sot|><|lang|><|task|><|notimestamps|> text <|eot|>).

Layout, float32 [n_text + 1, frames + 1]: alignment_matrix of the rows that predict each text token and end-of-text, then one column of token ids (exact in float32): row i’s text token, and on the last row the language token, which word splitting needs. Device ops only, so the walk never waits on the host.

Parameters:
Return type:

Tensor

mstar.model.whisper.components.alignment.dtw(cost)[source]#

Cheapest monotonic path from the top-left to the bottom-right of cost (rows x frames), moving right, down or diagonally. Returns the path’s row and frame indices, each of length len(path).

Parameters:

cost (ndarray)

Return type:

tuple[ndarray, ndarray]

mstar.model.whisper.components.alignment.median_filter(x, width=MEDIAN_FILTER_WIDTH)[source]#

Sliding median over the last dim (odd width), edges reflected.

Parameters:
Return type:

Tensor

mstar.model.whisper.components.alignment.merge_punctuation(words)[source]#

openai-whisper’s merge_punctuations: opening marks join the word after them (scanning right to left), then closing marks join the word before them. A word keeps its own timing; the mark’s is dropped.

Parameters:

words (list[dict])

Return type:

list[dict]

mstar.model.whisper.components.alignment.split_words(decode, token_ids, language=None)[source]#

Group text tokens into words, as openai-whisper’s split_to_word_tokens: first into pieces that decode to whole characters (a token ending inside a multi-byte character continues into the next), then, unless language is written without spaces, into words that begin at a leading space or a lone punctuation mark. decode(ids) -> str is the tokenizer’s plain decode, without clean-up, so the pieces’ texts add up to decode(token_ids).

Parameters:
Return type:

list[tuple[str, list[int]]]

mstar.model.whisper.components.alignment.token_start_frames(matrix)[source]#

The frame at which the DTW path first enters each row of matrix (a high value is a good alignment, so the path minimizes -matrix).

Parameters:

matrix (Tensor)

Return type:

ndarray

mstar.model.whisper.components.alignment.unpack_output(output)[source]#

alignment_output’s tensor -> (matrix, text token ids, language token id).

Parameters:

output (Tensor)

Return type:

tuple[Tensor, list[int], int]

mstar.model.whisper.components.alignment.warm_up()[source]#

Compile the DTW now, in the process that runs it (the API server’s data worker, at model load), rather than inside the first word-timestamp request.

Return type:

None

mstar.model.whisper.components.alignment.word_timings(matrix, text_tokens, decode, language=None)[source]#

Words with start / end seconds for one window.

matrix is alignment_matrix over the rows that predict text_tokens plus the final row that predicts end-of-text (len(text_tokens) + 1 rows); language is the transcript’s language code.

Parameters:
Return type:

list[dict]