bezzam HF Staff commited on
Commit
d1bc8db
·
verified ·
1 Parent(s): 7b4006a

Delete transformers

Browse files
transformers/requirements.txt DELETED
@@ -1,8 +0,0 @@
1
- torch
2
- transformers
3
- evaluate
4
- datasets
5
- librosa
6
- jiwer
7
- num2words
8
- peft
 
 
 
 
 
 
 
 
 
transformers/run_eval.py DELETED
@@ -1,448 +0,0 @@
1
- import argparse
2
- import os
3
- import re
4
- import torch
5
- from torch.nn.attention import sdpa_kernel, SDPBackend
6
- from transformers import AutoConfig, AutoModelForSpeechSeq2Seq, AutoModelForMultimodalLM, AutoModelForCTC, AutoProcessor, MODEL_FOR_MULTIMODAL_LM_MAPPING, MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING, MODEL_FOR_CTC_MAPPING, CompileConfig
7
- import evaluate
8
- from normalizer import data_utils
9
- from tqdm import tqdm
10
- import random
11
- import numpy as np
12
-
13
- wer_metric = evaluate.load("wer")
14
- torch.set_float32_matmul_precision('high')
15
-
16
-
17
- def remove_brackets(text):
18
- """
19
- Remove parentheses from text, replacing them with spaces.
20
-
21
- Some models (e.g. Cohere ASR) output parentheses that would cause the
22
- normalizer to delete the enclosed text entirely, leading to false
23
- deletion errors in the predictions.
24
- """
25
- text = text.replace("(", " ").replace(")", " ")
26
- text = re.sub(r'\s+', ' ', text)
27
- return text
28
-
29
-
30
- def main(args):
31
-
32
- # Set seed due to randomness in some models (e.g. VibeVoice's acoustic tokenizer sampling)
33
- seed = 42
34
- random.seed(seed)
35
- np.random.seed(seed)
36
- torch.manual_seed(seed)
37
- torch.cuda.manual_seed_all(seed)
38
- torch.backends.cudnn.deterministic = True
39
-
40
- torch_dtype = getattr(torch, args.dtype)
41
-
42
- config = AutoConfig.from_pretrained(args.model_id, revision=args.revision)
43
- if type(config) in MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING:
44
- cls_model = AutoModelForSpeechSeq2Seq
45
- elif type(config) in MODEL_FOR_MULTIMODAL_LM_MAPPING:
46
- cls_model = AutoModelForMultimodalLM
47
- elif type(config) in MODEL_FOR_CTC_MAPPING:
48
- cls_model = AutoModelForCTC
49
- else:
50
- raise ValueError(f"Model config of type {type(config)} not recognized in Transformers mappings.")
51
- is_ctc = cls_model == AutoModelForCTC
52
-
53
- if "vibevoice" in args.model_id.lower():
54
- model = cls_model.from_pretrained(
55
- args.model_id,
56
- dtype=torch_dtype,
57
- attn_implementation={
58
- "acoustic_tokenizer_encoder_config": "eager",
59
- "semantic_tokenizer_encoder_config": "eager",
60
- "text_config": "sdpa",
61
- }
62
- )
63
- else:
64
- model = cls_model.from_pretrained(
65
- args.model_id,
66
- dtype=torch_dtype,
67
- revision=args.revision,
68
- attn_implementation=args.attn_implementation,
69
- )
70
- model.to(args.device)
71
- model.eval()
72
- print(f"Model size: {sum(p.numel() for p in model.parameters()) / 1e9:.2f}B parameters")
73
- processor = AutoProcessor.from_pretrained(args.model_id, revision=args.revision)
74
- has_transcription_processor = hasattr(processor, "apply_transcription_request")
75
- is_cohere = "cohere" in args.model_id.lower() and "transcribe" in args.model_id.lower()
76
-
77
- # Optional prompt for audio language models, newer models should use `apply_transcription_request`
78
- text = None
79
- if "granite-speech-3.3" in args.model_id.lower():
80
- # create text prompt
81
- chat = [
82
- {
83
- "role": "system",
84
- "content": "Knowledge Cutoff Date: April 2024.\nToday's Date: December 19, 2024.\nYou are Granite, developed by IBM. You are a helpful AI assistant",
85
- },
86
- {
87
- "role": "user",
88
- "content": "<|audio|>can you transcribe the speech into a written format?",
89
- }
90
- ]
91
-
92
- text = processor.apply_chat_template(
93
- chat, tokenize=False, add_generation_prompt=True
94
- )
95
-
96
- # Extract sampling rate
97
- if hasattr(processor, "feature_extractor") and processor.feature_extractor is not None:
98
- sampling_rate = processor.feature_extractor.sampling_rate
99
- elif hasattr(processor, "audio_processor") and processor.audio_processor is not None:
100
- sampling_rate = processor.audio_processor.sampling_rate
101
- else:
102
- sampling_rate = 16_000
103
-
104
- # Set generate arguments
105
- if model.can_generate():
106
- gen_kwargs = {"max_new_tokens": args.max_new_tokens}
107
- if getattr(model.generation_config, "is_multilingual", False):
108
- gen_kwargs["language"] = "en"
109
- gen_kwargs["task"] = "transcribe"
110
- # Clear deprecated Whisper generation config fields to suppress warnings
111
- if hasattr(model.generation_config, "forced_decoder_ids"):
112
- model.generation_config.forced_decoder_ids = None
113
- if hasattr(model.generation_config, "suppress_tokens"):
114
- model.generation_config.suppress_tokens = []
115
- if hasattr(model.generation_config, "begin_suppress_tokens"):
116
- model.generation_config.begin_suppress_tokens = []
117
- if "granite-speech-3.3" in args.model_id.lower():
118
- gen_kwargs["repetition_penalty"] = 1.0
119
- elif args.max_new_tokens:
120
- raise ValueError("`max_new_tokens` should only be set for auto-regressive models, but got a CTC model.")
121
-
122
- if args.torch_compile is not None:
123
- if model.can_generate():
124
- gen_kwargs["compile_config"] = CompileConfig(mode=args.torch_compile, fullgraph=args.compile_fullgraph)
125
- # enable static k/v cache for autoregressive models
126
- model.generation_config.cache_implementation = "static"
127
- else:
128
- model = torch.compile(model, mode=args.torch_compile, fullgraph=args.compile_fullgraph)
129
-
130
- # Ensure warm-up runs when using torch.compile
131
- if args.warmup_steps is None or args.warmup_steps < 1:
132
- print("`--torch_compile` is enabled; forcing `--warmup_steps=10` to trigger compilation before timed runs.")
133
- args.warmup_steps = 10
134
-
135
- def benchmark(batch, min_new_tokens=None):
136
- # Load audio inputs
137
- audios = [audio["array"] for audio in batch["audio"]]
138
- minibatch_size = len(audios)
139
- sampling_rate = batch["audio"][0]["sampling_rate"]
140
- batch["audio_length_s"] = [len(audio) / sampling_rate for audio in audios]
141
- batch["audio_filepath"] = data_utils.extract_audio_filepaths_from_batch(batch, minibatch_size)
142
- if text is not None:
143
- texts=[text] * minibatch_size
144
- else:
145
- texts = None
146
-
147
- # START TIMING
148
- torch.cuda.synchronize(device=args.device)
149
- start_event = torch.cuda.Event(enable_timing=True)
150
- end_event = torch.cuda.Event(enable_timing=True)
151
- start_event.record()
152
-
153
- # 1. Pre-Processing
154
- # 1.1 Pad audios to max batch size if using torch compile to prevent re-compilations
155
- padding_size = None
156
- if minibatch_size != args.batch_size and args.torch_compile is not None:
157
- padding_size = args.batch_size - minibatch_size
158
- padding_audios = [audios[-1] for _ in range(padding_size)]
159
- audios.extend(padding_audios)
160
-
161
- if is_cohere:
162
- inputs = processor(
163
- audios,
164
- sampling_rate=sampling_rate,
165
- return_tensors="pt",
166
- language="en",
167
- punctuation=False,
168
- )
169
- elif has_transcription_processor:
170
- if "voxtral" in args.model_id.lower():
171
- inputs = processor.apply_transcription_request(
172
- language="en", # English for benchmark consistency
173
- audio=audios,
174
- model_id=args.model_id,
175
- sampling_rate=sampling_rate,
176
- format=["wav"] * len(audios),
177
- )
178
- else:
179
- inputs = processor.apply_transcription_request(audios)
180
- prompt_len = inputs["input_ids"].shape[1]
181
- elif texts is not None:
182
- inputs = processor(
183
- texts,
184
- audios,
185
- device=args.device, # Computation device; returned tensors are put on CPU
186
- return_tensors="pt",
187
- ).to(args.device)
188
- prompt_len = inputs["input_ids"].shape[1]
189
- elif not model.can_generate(): #or len(audios[0]) > processor.feature_extractor.n_samples:
190
- # 1.2 Either CTC pre-processing (normalize to mean 0, std 1), or long-form Whisper processing
191
- inputs = processor(
192
- audios,
193
- sampling_rate=sampling_rate,
194
- truncation=False,
195
- padding="longest",
196
- return_tensors="pt",
197
- return_attention_mask=True,
198
- )
199
- else:
200
- # 1.3 Standard Whisper processing: pad audios to 30-seconds and converted to log-mel
201
- if args.longform:
202
- inputs = processor(
203
- audios,
204
- sampling_rate=sampling_rate,
205
- return_tensors="pt",
206
- truncation=False,
207
- padding="longest",
208
- return_attention_mask=True,
209
- )
210
- else:
211
- inputs = processor(audios, sampling_rate=sampling_rate, return_tensors="pt", padding="longest", return_attention_mask=True, device=args.device)
212
-
213
- inputs = inputs.to(args.device, dtype=torch_dtype)
214
-
215
- # 2. Model Inference
216
- if args.torch_compile is not None:
217
- sdpa_backends = [SDPBackend.MATH]
218
- else:
219
- sdpa_backends = [SDPBackend.FLASH_ATTENTION, SDPBackend.EFFICIENT_ATTENTION, SDPBackend.MATH]
220
- with sdpa_kernel(sdpa_backends):
221
- if model.can_generate():
222
- # 2.1 Auto-regressive generation for LM-based models
223
- if args.longform:
224
- pred_ids = model.generate(**inputs, **gen_kwargs, return_timestamps=True)
225
- else:
226
- pred_ids = model.generate(**inputs, **gen_kwargs, min_new_tokens=min_new_tokens)
227
- else:
228
- # 2.2. Single forward pass for CTC
229
- with torch.no_grad():
230
- logits = model(**inputs).logits
231
- pred_ids = logits.argmax(-1)
232
-
233
- # 3. Post-processing
234
- # 3.1 Strip padded ids from predictions
235
- if padding_size is not None:
236
- pred_ids = pred_ids[:-padding_size, ...]
237
-
238
- # 3.2 Convert token ids to text transcription
239
- if is_cohere:
240
- audio_chunk_index = inputs.get("audio_chunk_index")
241
- pred_text = processor.decode(
242
- pred_ids, skip_special_tokens=True,
243
- audio_chunk_index=audio_chunk_index, language="en",
244
- )
245
- pred_text = [remove_brackets(t) for t in pred_text]
246
- elif "vibevoice" in args.model_id.lower():
247
- # VibeVoice: strip the input prompt tokens then use the model's own decode API
248
- generated_ids = pred_ids[:, prompt_len:]
249
- try:
250
- pred_text = processor.decode(generated_ids, return_format="transcription_only")
251
- except Exception as e:
252
- print(f"Batch decoding failed with error: {e}. Falling back to individual sample decoding.")
253
- pred_text = []
254
- for i, sample_ids in enumerate(generated_ids):
255
- try:
256
- decoded = processor.decode(sample_ids.unsqueeze(0), return_format="transcription_only")
257
- pred_text.append(decoded[0] if isinstance(decoded, list) else decoded)
258
- except Exception as sample_error:
259
- print(f"Sample {i} decoding failed with error: {sample_error}. Setting to empty transcript.")
260
- pred_text.append("")
261
- elif has_transcription_processor or texts is not None:
262
- # Strip input prompt tokens
263
- pred_text = processor.decode(pred_ids[:, prompt_len:], skip_special_tokens=True)
264
- elif is_ctc:
265
- # don't use skip_special_tokens as it collapses double letters
266
- pred_text = processor.batch_decode(pred_ids)
267
- else:
268
- pred_text = processor.decode(pred_ids, skip_special_tokens=True)
269
-
270
- # END TIMING
271
- end_event.record()
272
- torch.cuda.synchronize(device=args.device)
273
- runtime = start_event.elapsed_time(end_event) / 1000.0
274
-
275
- # normalize by minibatch size since we want the per-sample time
276
- batch["transcription_time_s"] = minibatch_size * [runtime / minibatch_size]
277
-
278
- # normalize transcriptions with English normalizer
279
- batch["predictions"] = [data_utils.normalizer(pred) for pred in pred_text]
280
- batch["references"] = batch["norm_text"]
281
- return batch
282
-
283
- if args.warmup_steps is not None:
284
- dataset = data_utils.load_data(args)
285
- dataset = data_utils.prepare_data(dataset, sampling_rate=sampling_rate)
286
-
287
- num_warmup_samples = args.warmup_steps * args.batch_size
288
- if args.streaming:
289
- warmup_dataset = dataset.take(num_warmup_samples)
290
- else:
291
- warmup_dataset = dataset.select(range(min(num_warmup_samples, len(dataset))))
292
- warmup_dataset = iter(warmup_dataset.map(benchmark, batch_size=args.batch_size, batched=True, fn_kwargs={"min_new_tokens": args.max_new_tokens}))
293
-
294
- for _ in tqdm(warmup_dataset, desc="Warming up..."):
295
- continue
296
-
297
- dataset = data_utils.load_data(args)
298
- if args.max_eval_samples is not None and args.max_eval_samples > 0:
299
- print(f"Subsampling dataset to first {args.max_eval_samples} samples!")
300
- if args.streaming:
301
- dataset = dataset.take(args.max_eval_samples)
302
- else:
303
- dataset = dataset.select(range(min(args.max_eval_samples, len(dataset))))
304
- dataset = data_utils.prepare_data(dataset, sampling_rate=sampling_rate)
305
-
306
- dataset = dataset.map(
307
- benchmark, batch_size=args.batch_size, batched=True, remove_columns=["audio"],
308
- )
309
-
310
- all_results = {
311
- "audio_length_s": [],
312
- "transcription_time_s": [],
313
- "predictions": [],
314
- "references": [],
315
- "audio_filepath": [],
316
- }
317
- result_iter = iter(dataset)
318
- for result in tqdm(result_iter, desc="Samples..."):
319
- for key in all_results:
320
- all_results[key].append(result[key])
321
-
322
- # Write manifest results (WER and RTFX)
323
- # Filtering of empty references is handled inside write_manifest.
324
- manifest_path = data_utils.write_manifest(
325
- all_results["references"],
326
- all_results["predictions"],
327
- args.model_id,
328
- args.dataset_path,
329
- args.dataset,
330
- args.split,
331
- audio_length=all_results["audio_length_s"],
332
- transcription_time=all_results["transcription_time_s"],
333
- audio_filepaths=all_results["audio_filepath"],
334
- )
335
- print("Results saved at path:", os.path.abspath(manifest_path))
336
-
337
- wer = wer_metric.compute(
338
- references=all_results["references"], predictions=all_results["predictions"]
339
- )
340
- wer = round(100 * wer, 2)
341
- rtfx = round(sum(all_results["audio_length_s"]) / sum(all_results["transcription_time_s"]), 2)
342
- print("WER:", wer, "%", "RTFx:", rtfx)
343
-
344
-
345
- if __name__ == "__main__":
346
- parser = argparse.ArgumentParser()
347
-
348
- parser.add_argument(
349
- "--model_id",
350
- type=str,
351
- required=True,
352
- help="Model identifier. Should be loadable with 🤗 Transformers",
353
- )
354
- parser.add_argument(
355
- "--dataset_path",
356
- type=str,
357
- default="hf-audio/open-asr-leaderboard",
358
- help="Dataset path. By default, it is `hf-audio/open-asr-leaderboard`",
359
- )
360
- parser.add_argument(
361
- "--dataset",
362
- type=str,
363
- required=True,
364
- help="Dataset name. *E.g.* `'librispeech_asr` for the LibriSpeech ASR dataset, or `'common_voice'` for Common Voice. The full list of dataset names "
365
- "can be found at `https://huggingface.co/datasets/hf-audio/open-asr-leaderboard`",
366
- )
367
- parser.add_argument(
368
- "--split",
369
- type=str,
370
- default="test",
371
- help="Split of the dataset. *E.g.* `'validation`' for the dev split, or `'test'` for the test split.",
372
- )
373
- parser.add_argument(
374
- "--device",
375
- type=int,
376
- default=-1,
377
- help="The device to run the pipeline on. -1 for CPU (default), 0 for the first GPU and so on.",
378
- )
379
- parser.add_argument(
380
- "--batch_size",
381
- type=int,
382
- default=16,
383
- help="Number of samples to go through each streamed batch.",
384
- )
385
- parser.add_argument(
386
- "--max_eval_samples",
387
- type=int,
388
- default=None,
389
- help="Number of samples to be evaluated. Put a lower number e.g. 64 for testing this script.",
390
- )
391
- parser.add_argument(
392
- "--streaming",
393
- action="store_true",
394
- help="Stream the dataset lazily over the network instead of downloading it in full before the evaluation. Off by default for reproducible benchmark timings.",
395
- )
396
- parser.add_argument(
397
- "--max_new_tokens",
398
- type=int,
399
- default=None,
400
- help="Maximum number of tokens to generate (for auto-regressive models).",
401
- )
402
- parser.add_argument(
403
- "--longform",
404
- action="store_true",
405
- help="Whether to use longform mode.",
406
- )
407
- parser.add_argument(
408
- "--torch_compile",
409
- type=str,
410
- default=None,
411
- help="Mode for torch compiling model forward pass. Can be either 'default', 'reduce-overhead', 'max-autotune' or 'max-autotune-no-cudagraphs'.",
412
- )
413
- parser.add_argument(
414
- "--compile_fullgraph",
415
- action="store_true",
416
- help="Whether to do full graph compilation.",
417
- )
418
- parser.add_argument(
419
- "--dtype",
420
- type=str,
421
- default="bfloat16",
422
- help="The dtype to use for model loading and inference. E.g. 'bfloat16', 'float16', 'float32'.",
423
- )
424
- parser.add_argument(
425
- "--attn_implementation",
426
- type=str,
427
- default="sdpa",
428
- help="Attention implementation to use for model loading (e.g. 'sdpa', 'eager', 'flash_attention_2').",
429
- )
430
- parser.add_argument(
431
- "--warmup_steps",
432
- type=int,
433
- default=10,
434
- help="Number of warm-up steps to run before launching the timed runs.",
435
- )
436
- parser.add_argument(
437
- "--revision",
438
- type=str,
439
- default=None,
440
- help="Model revision to use (e.g. 'refs/pr/11' for a PR branch). Defaults to the main branch.",
441
- )
442
- args = parser.parse_args()
443
-
444
- print("*" * 100)
445
- print(f"Evaluating {args.model_id} on {args.dataset_path} / {args.dataset} / {args.split}")
446
- print("*" * 100)
447
-
448
- main(args)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
transformers/run_eval_ml.py DELETED
@@ -1,389 +0,0 @@
1
- import argparse
2
- import os
3
- import torch
4
- from torch.nn.attention import sdpa_kernel, SDPBackend
5
- from transformers import AutoConfig, AutoModelForSpeechSeq2Seq, AutoModelForMultimodalLM, AutoModelForCTC, AutoProcessor, MODEL_FOR_MULTIMODAL_LM_MAPPING, MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING, MODEL_FOR_CTC_MAPPING, CompileConfig
6
- import evaluate
7
- from normalizer import data_utils
8
- from normalizer.eval_utils import normalize_compound_pairs
9
- from tqdm import tqdm
10
- from datasets import load_dataset, Audio
11
- import random
12
- import numpy as np
13
-
14
- wer_metric = evaluate.load("wer")
15
- torch.set_float32_matmul_precision('high')
16
-
17
-
18
- def main(args):
19
-
20
- # Set seed for reproducibility
21
- seed = 42
22
- random.seed(seed)
23
- np.random.seed(seed)
24
- torch.manual_seed(seed)
25
- torch.cuda.manual_seed_all(seed)
26
- torch.backends.cudnn.deterministic = True
27
-
28
- torch_dtype = getattr(torch, args.dtype)
29
-
30
- config = AutoConfig.from_pretrained(args.model_id)
31
- if type(config) in MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING:
32
- cls_model = AutoModelForSpeechSeq2Seq
33
- elif type(config) in MODEL_FOR_MULTIMODAL_LM_MAPPING:
34
- cls_model = AutoModelForMultimodalLM
35
- elif type(config) in MODEL_FOR_CTC_MAPPING:
36
- cls_model = AutoModelForCTC
37
- else:
38
- raise ValueError(f"Model config of type {type(config)} not recognized in Transformers mappings.")
39
- is_ctc = cls_model == AutoModelForCTC
40
-
41
- model = cls_model.from_pretrained(
42
- args.model_id,
43
- dtype=torch_dtype,
44
- attn_implementation=args.attn_implementation,
45
- )
46
- model.to(args.device)
47
- model.eval()
48
- print(f"Model size: {sum(p.numel() for p in model.parameters()) / 1e9:.2f}B parameters")
49
- processor = AutoProcessor.from_pretrained(args.model_id)
50
- has_transcription_processor = hasattr(processor, "apply_transcription_request")
51
-
52
- # Extract sampling rate from processor
53
- if hasattr(processor, "feature_extractor") and processor.feature_extractor is not None:
54
- sampling_rate = processor.feature_extractor.sampling_rate
55
- elif hasattr(processor, "audio_processor") and processor.audio_processor is not None:
56
- sampling_rate = processor.audio_processor.sampling_rate
57
- else:
58
- sampling_rate = 16_000
59
-
60
- # Set generate arguments (only for auto-regressive models)
61
- if model.can_generate():
62
- gen_kwargs = {}
63
- if args.max_new_tokens is not None:
64
- gen_kwargs["max_new_tokens"] = args.max_new_tokens
65
-
66
- # For multilingual models, set task to transcribe and pass language (None = auto-detect)
67
- if getattr(model.generation_config, "is_multilingual", False):
68
- gen_kwargs["task"] = "transcribe"
69
- if args.language is not None:
70
- gen_kwargs["language"] = args.language
71
- elif args.max_new_tokens:
72
- raise ValueError("`max_new_tokens` should only be set for auto-regressive models, but got a CTC model.")
73
-
74
- CONFIG_NAME = args.config_name
75
- SPLIT_NAME = args.split
76
-
77
- # Determine language for normalization: use --language if provided, otherwise extract from config_name
78
- if args.language is not None:
79
- norm_language = args.language
80
- else:
81
- try:
82
- norm_language = CONFIG_NAME.split("_", 1)[1]
83
- except IndexError:
84
- norm_language = "en"
85
- print(f"Language not specified, extracted '{norm_language}' from config_name '{CONFIG_NAME}'")
86
-
87
- if args.torch_compile is not None:
88
- if model.can_generate():
89
- gen_kwargs["compile_config"] = CompileConfig(mode=args.torch_compile, fullgraph=args.compile_fullgraph)
90
- model.generation_config.cache_implementation = "static"
91
- else:
92
- model = torch.compile(model, mode=args.torch_compile, fullgraph=args.compile_fullgraph)
93
-
94
- # Ensure warm-up runs when using torch.compile
95
- if args.warmup_steps is None or args.warmup_steps < 1:
96
- print("`--torch_compile` is enabled; forcing `--warmup_steps=10` to trigger compilation before timed runs.")
97
- args.warmup_steps = 10
98
-
99
- # Load dataset
100
- print(f"Loading dataset: {args.dataset} with config: {CONFIG_NAME}")
101
- dataset = load_dataset(
102
- args.dataset,
103
- CONFIG_NAME,
104
- split=SPLIT_NAME,
105
- streaming=args.streaming,
106
- token=True,
107
- )
108
- dataset = dataset.cast_column("audio", Audio(sampling_rate=sampling_rate))
109
-
110
- if args.max_eval_samples is not None and args.max_eval_samples > 0:
111
- print(f"Subsampling dataset to first {args.max_eval_samples} samples!")
112
- if args.streaming:
113
- dataset = dataset.take(args.max_eval_samples)
114
- else:
115
- dataset = dataset.select(range(min(args.max_eval_samples, len(dataset))))
116
-
117
- def benchmark(batch, min_new_tokens=None):
118
- audios = [audio["array"] for audio in batch["audio"]]
119
- minibatch_size = len(audios)
120
- sampling_rate = batch["audio"][0]["sampling_rate"]
121
- batch["audio_length_s"] = [len(audio) / sampling_rate for audio in audios]
122
- batch["audio_filepath"] = data_utils.extract_audio_filepaths_from_batch(batch, minibatch_size)
123
-
124
- # START TIMING
125
- torch.cuda.synchronize(device=args.device)
126
- start_event = torch.cuda.Event(enable_timing=True)
127
- end_event = torch.cuda.Event(enable_timing=True)
128
- start_event.record()
129
-
130
- # 1. Pre-Processing
131
- # Pad audios to max batch size if using torch compile to prevent re-compilations
132
- padding_size = None
133
- if minibatch_size != args.batch_size and args.torch_compile is not None:
134
- padding_size = args.batch_size - minibatch_size
135
- padding_audios = [audios[-1] for _ in range(padding_size)]
136
- audios.extend(padding_audios)
137
-
138
- if has_transcription_processor:
139
- if "voxtral" in args.model_id.lower():
140
- inputs = processor.apply_transcription_request(
141
- language=args.language, # None = auto-detect
142
- audio=audios,
143
- model_id=args.model_id,
144
- sampling_rate=sampling_rate,
145
- format=["wav"] * len(audios),
146
- )
147
- else:
148
- inputs = processor.apply_transcription_request(audios)
149
- prompt_len = inputs["input_ids"].shape[1]
150
- elif not model.can_generate():
151
- # CTC pre-processing: normalize to mean 0, std 1
152
- inputs = processor(
153
- audios,
154
- sampling_rate=sampling_rate,
155
- truncation=False,
156
- padding="longest",
157
- return_tensors="pt",
158
- return_attention_mask=True,
159
- )
160
- else:
161
- # Standard Whisper processing: pad audios to 30-seconds and convert to log-mel
162
- inputs = processor(
163
- audios,
164
- sampling_rate=sampling_rate,
165
- return_tensors="pt",
166
- padding="longest",
167
- return_attention_mask=True,
168
- device=args.device,
169
- )
170
-
171
- inputs = inputs.to(args.device, dtype=torch_dtype)
172
-
173
- # 2. Model Inference
174
- if args.torch_compile is not None:
175
- sdpa_backends = [SDPBackend.MATH]
176
- else:
177
- sdpa_backends = [SDPBackend.FLASH_ATTENTION, SDPBackend.EFFICIENT_ATTENTION, SDPBackend.MATH]
178
- with sdpa_kernel(sdpa_backends):
179
- if model.can_generate():
180
- pred_ids = model.generate(**inputs, **gen_kwargs, min_new_tokens=min_new_tokens)
181
- else:
182
- # Single forward pass for CTC
183
- with torch.no_grad():
184
- logits = model(**inputs).logits
185
- pred_ids = logits.argmax(-1)
186
-
187
- # 3. Post-processing
188
- # Strip padded ids from predictions
189
- if padding_size is not None:
190
- pred_ids = pred_ids[:-padding_size, ...]
191
-
192
- # Convert token ids to text transcription
193
- if has_transcription_processor:
194
- pred_text = processor.batch_decode(pred_ids[:, prompt_len:], skip_special_tokens=True)
195
- elif is_ctc:
196
- # don't use skip_special_tokens as it collapses double letters
197
- pred_text = processor.batch_decode(pred_ids)
198
- else:
199
- pred_text = processor.batch_decode(pred_ids, skip_special_tokens=True)
200
-
201
- # END TIMING
202
- end_event.record()
203
- torch.cuda.synchronize(device=args.device)
204
- runtime = start_event.elapsed_time(end_event) / 1000.0
205
-
206
- batch["transcription_time_s"] = minibatch_size * [runtime / minibatch_size]
207
-
208
- # Normalize with multilingual normalizer
209
- batch["predictions"] = [data_utils.ml_normalizer(pred, lang=norm_language) for pred in pred_text]
210
- batch["references"] = [data_utils.ml_normalizer(ref, lang=norm_language) for ref in batch["text"]]
211
-
212
- return batch
213
-
214
- if args.warmup_steps is not None and args.warmup_steps > 0:
215
- print(f"Running {args.warmup_steps} warmup steps...")
216
- num_warmup_samples = args.warmup_steps * args.batch_size
217
- if args.streaming:
218
- warmup_dataset = dataset.take(num_warmup_samples)
219
- else:
220
- warmup_dataset = dataset.select(range(min(num_warmup_samples, len(dataset))))
221
- warmup_dataset = iter(warmup_dataset.map(
222
- benchmark, batch_size=args.batch_size, batched=True,
223
- fn_kwargs={"min_new_tokens": args.max_new_tokens}
224
- ))
225
- for _ in tqdm(warmup_dataset, desc="Warming up..."):
226
- continue
227
-
228
- # Reload dataset for actual evaluation (reset streaming pointer)
229
- dataset = load_dataset(
230
- args.dataset,
231
- CONFIG_NAME,
232
- split=SPLIT_NAME,
233
- streaming=args.streaming,
234
- token=True,
235
- )
236
- dataset = dataset.cast_column("audio", Audio(sampling_rate=sampling_rate))
237
-
238
- if args.max_eval_samples is not None and args.max_eval_samples > 0:
239
- if args.streaming:
240
- dataset = dataset.take(args.max_eval_samples)
241
- else:
242
- dataset = dataset.select(range(min(args.max_eval_samples, len(dataset))))
243
-
244
- dataset = dataset.map(
245
- benchmark, batch_size=args.batch_size, batched=True, remove_columns=["audio"],
246
- )
247
-
248
- all_results = {
249
- "audio_length_s": [],
250
- "transcription_time_s": [],
251
- "predictions": [],
252
- "references": [],
253
- "audio_filepath": [],
254
- }
255
-
256
- result_iter = iter(dataset)
257
- for result in tqdm(result_iter, desc="Samples..."):
258
- for key in all_results:
259
- all_results[key].append(result[key])
260
-
261
- # Filter empty references (consistent with English pipeline)
262
- filtered = [
263
- (ref, pred, dur, time_s, fpath)
264
- for ref, pred, dur, time_s, fpath in zip(
265
- all_results["references"], all_results["predictions"],
266
- all_results["audio_length_s"], all_results["transcription_time_s"],
267
- all_results["audio_filepath"]
268
- )
269
- if data_utils.is_target_text_in_range(ref)
270
- ]
271
- if filtered:
272
- all_results["references"], all_results["predictions"], all_results["audio_length_s"], all_results["transcription_time_s"], all_results["audio_filepath"] = zip(*filtered)
273
- all_results = {k: list(v) for k, v in all_results.items()}
274
-
275
- # Write manifest results (WER and RTFX)
276
- manifest_path = data_utils.write_manifest(
277
- all_results["references"],
278
- all_results["predictions"],
279
- args.model_id,
280
- args.dataset,
281
- CONFIG_NAME,
282
- args.split,
283
- audio_length=all_results["audio_length_s"],
284
- transcription_time=all_results["transcription_time_s"],
285
- audio_filepaths=all_results["audio_filepath"],
286
- )
287
- print("Results saved at path:", os.path.abspath(manifest_path))
288
-
289
- wer_refs, wer_preds = normalize_compound_pairs(all_results["references"], all_results["predictions"])
290
- wer = wer_metric.compute(references=wer_refs, predictions=wer_preds)
291
- wer = round(100 * wer, 2)
292
- rtfx = round(sum(all_results["audio_length_s"]) / sum(all_results["transcription_time_s"]), 2)
293
- print("WER:", wer, "%", "RTFx:", rtfx)
294
-
295
-
296
- if __name__ == "__main__":
297
- parser = argparse.ArgumentParser()
298
-
299
- parser.add_argument(
300
- "--model_id",
301
- type=str,
302
- required=True,
303
- help="Model identifier. Should be loadable with Transformers",
304
- )
305
- parser.add_argument(
306
- "--dataset",
307
- type=str,
308
- required=True,
309
- help="Dataset name. E.g. 'nithinraok/asr-leaderboard-datasets'",
310
- )
311
- parser.add_argument(
312
- "--config_name",
313
- type=str,
314
- required=True,
315
- help="Config name for the dataset. E.g. 'fleurs_de' for German FLEURS.",
316
- )
317
- parser.add_argument(
318
- "--language",
319
- type=str,
320
- default=None,
321
- help="Language code, e.g. 'de' for German. If not set, the model will auto-detect the language.",
322
- )
323
- parser.add_argument(
324
- "--split",
325
- type=str,
326
- default="test",
327
- help="Split of the dataset.",
328
- )
329
- parser.add_argument(
330
- "--device",
331
- type=int,
332
- default=-1,
333
- help="The device to run the pipeline on. -1 for CPU (default), 0 for the first GPU and so on.",
334
- )
335
- parser.add_argument(
336
- "--batch_size",
337
- type=int,
338
- default=64,
339
- help="Number of samples to go through each batch.",
340
- )
341
- parser.add_argument(
342
- "--max_eval_samples",
343
- type=int,
344
- default=None,
345
- help="Number of samples to be evaluated. Put a lower number e.g. 64 for testing this script.",
346
- )
347
- parser.add_argument(
348
- "--streaming",
349
- action="store_true",
350
- help="Stream the dataset lazily over the network instead of downloading it in full before the evaluation. Off by default for reproducible benchmark timings.",
351
- )
352
- parser.add_argument(
353
- "--max_new_tokens",
354
- type=int,
355
- default=None,
356
- help="Maximum number of tokens to generate.",
357
- )
358
- parser.add_argument(
359
- "--torch_compile",
360
- type=str,
361
- default=None,
362
- help="Mode for torch compiling model forward pass. Can be either 'default', 'reduce-overhead', 'max-autotune' or 'max-autotune-no-cudagraphs'.",
363
- )
364
- parser.add_argument(
365
- "--compile_fullgraph",
366
- action="store_true",
367
- help="Whether to do full graph compilation.",
368
- )
369
- parser.add_argument(
370
- "--dtype",
371
- type=str,
372
- default="bfloat16",
373
- help="The dtype to use for model loading and inference. E.g. 'bfloat16', 'float16', 'float32'.",
374
- )
375
- parser.add_argument(
376
- "--attn_implementation",
377
- type=str,
378
- default="sdpa",
379
- help="Attention implementation to use for model loading (e.g. 'sdpa', 'eager', 'flash_attention_2').",
380
- )
381
- parser.add_argument(
382
- "--warmup_steps",
383
- type=int,
384
- default=10,
385
- help="Number of warm-up steps to run before launching the timed runs.",
386
- )
387
- args = parser.parse_args()
388
-
389
- main(args)