in ctranslate2/run_eval.py [0:0]
def main(args) -> None:
"""Main function to run evaluation on a dataset."""
asr_model = WhisperModel(
model_size_or_path=args.model_id,
compute_type="float16",
device="cuda",
device_index=args.device
)
def benchmark(batch):
start_time = time.time()
segments, _ = asr_model.transcribe(batch["audio"]["array"], language="en")
outputs = [segment._asdict() for segment in segments]
batch["transcription_time_s"] = time.time() - start_time
batch["predictions"] = data_utils.normalizer("".join([segment["text"] for segment in outputs])).strip()
batch["references"] = batch["norm_text"]
return batch
if args.warmup_steps is not None:
dataset = data_utils.load_data(args)
dataset = data_utils.prepare_data(dataset)
if args.streaming:
warmup_dataset = dataset.take(args.warmup_steps)
else:
warmup_dataset = dataset.select(range(min(args.warmup_steps, len(dataset))))
warmup_dataset = iter(warmup_dataset.map(benchmark, remove_columns=["audio"]))
for _ in tqdm(warmup_dataset, desc="Warming up..."):
continue
dataset = data_utils.load_data(args)
if args.max_eval_samples is not None and args.max_eval_samples > 0:
print(f"Subsampling dataset to first {args.max_eval_samples} samples!")
if args.streaming:
dataset = dataset.take(args.max_eval_samples)
else:
dataset = dataset.select(range(min(args.max_eval_samples, len(dataset))))
dataset = data_utils.prepare_data(dataset)
dataset = dataset.map(benchmark, remove_columns=["audio"])
all_results = {
"audio_length_s": [],
"transcription_time_s": [],
"predictions": [],
"references": [],
}
result_iter = iter(dataset)
for result in tqdm(result_iter, desc="Samples..."):
for key in all_results:
all_results[key].append(result[key])
# Write manifest results (WER and RTFX)
manifest_path = data_utils.write_manifest(
all_results["references"],
all_results["predictions"],
args.model_id,
args.dataset_path,
args.dataset,
args.split,
audio_length=all_results["audio_length_s"],
transcription_time=all_results["transcription_time_s"],
)
print("Results saved at path:", os.path.abspath(manifest_path))
wer = wer_metric.compute(
references=all_results["references"], predictions=all_results["predictions"]
)
wer = round(100 * wer, 2)
rtfx = round(sum(all_results["audio_length_s"]) / sum(all_results["transcription_time_s"]), 2)
print("WER:", wer, "%", "RTFx:", rtfx)