File size: 9,413 Bytes
3c3b55b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a4ed886
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
3c3b55b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a4ed886
3c3b55b
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
# Whisper: speech recognition with Transformers

This page is the Transformers contract of the Whisper MIC and also a usage example. [doc-builder](https://github.com/huggingface/doc-builder/blob/main/docs/runnable-code-blocks.md) renders the `runnable` blocks below as documentation, and its pytest plugin runs them as written against `openai/whisper-tiny`. The MIC runner runs the same blocks against each slim fixture and each checkpoint in [`integration.json`](../../integration.json). Lines marked `# doc-builder: hide` only matter there: they swap in the pinned checkpoint or variant and print the result record.

The audio is synthetic, so the example runs offline: four one-second clips, each a distinct pair of tones. A real checkpoint transcribes them as noise; the slim fixture is trained to name them.

## Transcribe a batch

```py runnable:transcribe
import math

import torch
from transformers import WhisperForConditionalGeneration, WhisperProcessor, set_seed
import json  # doc-builder: hide
import os  # doc-builder: hide

model_id = "openai/whisper-tiny"
model_id = os.environ.get("MIC_MODEL", model_id)  # doc-builder: hide
variant = json.loads(os.environ.get("MIC_MODEL_KWARGS", "{}"))  # doc-builder: hide
device = "cuda" if torch.cuda.is_available() else "cpu"
set_seed(0)  # seeds torch, NumPy, and Python: loading, dropout, and data augmentation draw random numbers

processor = WhisperProcessor.from_pretrained(model_id)
model = WhisperForConditionalGeneration.from_pretrained(model_id).to(device)
if variant:  # doc-builder: hide
    model = WhisperForConditionalGeneration.from_pretrained(model_id, **variant).to(device)

sampling_rate = 16000
half = torch.arange(sampling_rate // 2) / sampling_rate
tones = [(220, 880), (440, 1320), (660, 330), (990, 550)]
clips = [torch.cat([torch.sin(2 * math.pi * a * half), torch.sin(2 * math.pi * b * half)]).mul(0.5).numpy() for a, b in tones]
inputs = processor(clips, sampling_rate=sampling_rate, return_tensors="pt").to(device, model.dtype)  # floating-point inputs follow the model's dtype

generated = model.generate(inputs.input_features, language="en", task="transcribe", max_new_tokens=8)
print(processor.batch_decode(generated, skip_special_tokens=True))
```

## Transcribe with timestamps

Whisper can predict segment timestamps; decoding with `decode_with_timestamps=True` keeps them in the text. Integrations also read the feature extractor's configuration without building a processor.

```py runnable:transcribe:3
from transformers.models.auto.feature_extraction_auto import get_feature_extractor_config

timestamped = model.generate(inputs.input_features, language="en", task="transcribe", return_timestamps=True, max_new_tokens=8)
timestamped_text = processor.batch_decode(timestamped, decode_with_timestamps=True)
print(timestamped_text)

feature_config = get_feature_extractor_config(model_id)
print(sorted(feature_config))
```

## Limit the transcript length

For Whisper, `max_length` counts transcript tokens: `generate` adds the forced start, language, and task tokens on top.

```py runnable:transcribe:3
limited = model.generate(inputs.input_features, language="en", task="transcribe", max_length=1)
print(processor.batch_decode(limited, skip_special_tokens=True))
```

## Speculative decoding

A smaller decoder-only Whisper can draft tokens that the full model verifies, as [Distil-Whisper](https://huggingface.co/distil-whisper/distil-large-v3) does. Here the checkpoint's own decoder, loaded as `WhisperForCausalLM`, is the assistant; the transcript must not change. Assisted generation takes one clip at a time.

```py runnable:transcribe:3
from transformers import WhisperForCausalLM

assistant = WhisperForCausalLM.from_pretrained(model_id).to(device)
if variant:  # doc-builder: hide
    assistant = WhisperForCausalLM.from_pretrained(model_id, **variant).to(device)
assisted = [
    model.generate(features[None], language="en", task="transcribe", max_new_tokens=8, assistant_model=assistant)[0].tolist()
    for features in inputs.input_features
]
print(processor.batch_decode(assisted, skip_special_tokens=True))
```

## Transcribe long audio in a batch

Audio longer than 30 seconds is transcribed segment by segment. Keep the full audio (`truncation=False`), pad the batch to its longest clip, and pass the attention mask; `condition_on_prev_tokens=True` prompts each segment with the previous one's text.

```py runnable:transcribe:3
long_clips = [
    torch.cat([torch.from_numpy(clips[(start + k) % len(clips)]) for k in range(seconds)]).numpy()
    for start, seconds in ((0, 36), (1, 32))
]
long_inputs = processor(
    long_clips, sampling_rate=sampling_rate, return_tensors="pt", truncation=False, padding="longest", return_attention_mask=True
).to(device, model.dtype)
longform = model.generate(
    **long_inputs, language="en", task="transcribe", condition_on_prev_tokens=True, return_timestamps=True,
    temperature=0.0, max_new_tokens=8,
)
print(processor.batch_decode(longform, skip_special_tokens=True))
```

## Continue from transcripts of different lengths

Previous-segment transcripts can have different lengths across a batch. Whisper left-pads the shorter prompt before decoding the next segment. This reproducible example constrains the first segment to one word in the first row and three words in the second; the remaining audio is transcribed freely. The custom processor only chooses the first segment's tokens.

```py runnable:transcribe:3
from transformers import LogitsProcessor

tokenizer = processor.tokenizer
timestamp_begin = tokenizer.convert_tokens_to_ids("<|0.00|>")
previous_token = tokenizer.convert_tokens_to_ids("<|startofprev|>")
first_segment_tokens = [
    [timestamp_begin] + tokenizer.encode(text, add_special_tokens=False) + [timestamp_begin + 1500, model.config.eos_token_id]
    for text in (" red", " green blue yellow")
]

class FirstSegmentTranscript(LogitsProcessor):
    def __call__(self, input_ids, scores):
        if any(previous_token in row.tolist() for row in input_ids):
            return scores
        # The initial prefix is start-of-transcript, English, transcribe.
        step = input_ids.shape[-1] - 3
        constrained = torch.full_like(scores, -float("inf"))
        for row, tokens in enumerate(first_segment_tokens):
            token = tokens[step] if step < len(tokens) else model.config.eos_token_id
            constrained[row, token] = 0
        return constrained

continuation_inputs = processor(
    [clip.repeat(36) for clip in clips[:2]], sampling_rate=sampling_rate, return_tensors="pt",
    truncation=False, padding="longest", return_attention_mask=True,
).to(device, model.dtype)
continued = model.generate(
    **continuation_inputs, language="en", task="transcribe", condition_on_prev_tokens=True,
    return_timestamps=True, temperature=0.0, max_new_tokens=12,
    logits_processor=[FirstSegmentTranscript()],
)
print(processor.batch_decode(continued, skip_special_tokens=True))
```

## One training step

```py runnable:transcribe:3
transcripts = [" red", " green", " blue", " yellow"]
processor.tokenizer.set_prefix_tokens(language="english", task="transcribe", predict_timestamps=False)
targets = processor.tokenizer(transcripts, return_tensors="pt", padding=True).to(device)
decoder_input_ids = targets.input_ids[:, :-1]
labels = targets.input_ids[:, 1:].masked_fill(targets.attention_mask[:, 1:] == 0, -100)

with torch.no_grad():  # doc-builder: hide
    # Teacher-forced log-probabilities of the reference transcripts.
    logprobs = model(input_features=inputs.input_features, decoder_input_ids=decoder_input_ids).logits.float().log_softmax(-1)
    label_logprobs = logprobs.gather(-1, labels.clamp(min=0)[..., None])[..., 0].masked_fill(labels == -100, 0.0)

model.train()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
loss = model(input_features=inputs.input_features, decoder_input_ids=decoder_input_ids, labels=labels).loss
loss.backward()
before = {name: p.detach().clone() for name, p in model.named_parameters()}  # doc-builder: hide
optimizer.step()
print(f"loss: {loss.item():.4f}")

if os.environ.get("MIC_RUNNER"):  # doc-builder: hide
    named = dict(model.named_parameters())
    checked = ("model.encoder.conv1.weight", "model.decoder.layers.0.encoder_attn.q_proj.weight")
    result = {
        "runtime": {
            "device": model.device.type,
            "dtype": str(model.dtype).removeprefix("torch."),
            "attn_implementation": model.config._attn_implementation,
        },
        "generated_tokens": generated.tolist(),
        "generated_text": processor.batch_decode(generated, skip_special_tokens=True),
        "label_logprobs": label_logprobs.tolist(),
        "timestamped_tokens": timestamped.tolist(),
        "timestamped_text": timestamped_text,
        "feature_extractor_config_keys": sorted(feature_config),
        "max_length_tokens": limited.tolist(),
        "assisted_tokens": assisted,
        "longform_tokens": longform.tolist(),
        "continued_longform_tokens": continued.tolist(),
        "training": {
            "loss": loss.item(),
            "gradients": {n: bool(named[n].grad is not None and torch.count_nonzero(named[n].grad) > 0 and torch.isfinite(named[n].grad).all()) for n in checked},
            "updated": any(not torch.equal(p, before[n]) for n, p in named.items()),
        },
    }
    print(json.dumps(result))
```