whisper 推理脚本

import argparse import functools import gc import os import evaluate # Hugging Face评估库 import numpy as np import torch from torch.utils.data import DataLoader from tqdm import tqdm # 进度条显示 from transformers import WhisperForConditionalGeneration, WhisperProcessor # Whisper模型核心组件 # 自定义工具模块 from utils.data_utils import DataCollatorSpeechSeq2SeqWithPadding, remove_punctuation, to_simple from utils.reader import CustomDataset # 自定义数据集加载器 from utils.utils import print_arguments, add_arguments # 参数处理工具 # 创建命令行参数解析器 parser argparse.ArgumentParser(description__doc__) add_arg functools.partial(add_arguments, argparserparser) # 简化参数添加 # 添加评估所需的各种参数 add_arg(test_data, typestr, defaultdataset/test.json, help测试集的路径JSON格式) add_arg(model_path, typestr, defaultmodels/whisper-tiny-finetune, help合并模型的路径或者是huggingface上模型的名称) add_arg(batch_size, typeint, default16, help评估的batch size每次处理的音频数量) add_arg(num_workers, typeint, default8, help读取数据的线程数量多进程加载) add_arg(language, typestr, defaultChinese, help设置语言可全称也可简写如果为None则评估的是多语言) add_arg(remove_pun, typebool, defaultTrue, help是否移除标点符号影响评估指标计算) add_arg(to_simple, typebool, defaultFalse, help是否转为简体中文仅中文有效) add_arg(timestamps, typebool, defaultFalse, help评估时是否使用时间戳数据Whisper特有) add_arg(min_audio_len, typefloat, default0.5, help最小的音频长度单位秒过滤太短音频) add_arg(max_audio_len, typefloat, default30, help最大的音频长度单位秒过滤太长音频) add_arg(local_files_only, typebool, defaultTrue, help是否只在本地加载模型不尝试下载) add_arg(using_language_label, typebool, defaultFalse, help是否使用语言标签多语言模型) add_arg(decoder_start_index, typeint, default4, helpdecoder_start_index控制生成起始位置) add_arg(task, typestr, defaulttranscribe, choices[transcribe, translate], help模型的任务转录或翻译) add_arg(metric, typestr, defaultcer, choices[cer, wer], help评估方式字符错误率或词错误率) args parser.parse_args() print_arguments(args) # 打印所有参数配置 # 验证模型路径合法性 assert openai os.path.dirname(args.model_path) or os.path.exists(args.model_path), \ f模型文件{args.model_path}不存在请检查是否已经成功合并模型或者是否为huggingface存在模型 # 指定GPU设备使用第二块GPU device_map cuda:1 def main(): # 模型加载与配置 # 初始化Whisper处理器包含特征提取器分词器 processor WhisperProcessor.from_pretrained( args.model_path, languageargs.language, # 设置目标语言 taskargs.task, # 任务类型转录/翻译 no_timestampsnot args.timestamps, # 是否禁用时间戳 local_files_onlyargs.local_files_only # 仅使用本地文件 ) # 加载预训练模型 model WhisperForConditionalGeneration.from_pretrained( args.model_path, local_files_onlyargs.local_files_only ) # 关键配置调整解决生成问题 print(forced_decoder_ids: , model.config.forced_decoder_ids) # 打印原始强制解码配置 print(model.config.suppress_tokens: , model.config.suppress_tokens) # 打印抑制词配置 # 禁用强制解码规则允许模型自由生成 model.config.forced_decoder_ids None # 清空抑制词列表避免某些字符无法生成 model.config.suppress_tokens [] # 设置模型为评估模式 model.eval() # 将模型移动到指定GPU model model.to(device_map) # 数据准备 # 创建自定义测试数据集 test_dataset CustomDataset( data_list_pathargs.test_data, # 测试集路径 processorprocessor, # 处理器用于特征提取 timestampsargs.timestamps, # 是否包含时间戳 min_durationargs.min_audio_len, # 最小音频时长 max_durationargs.max_audio_len # 最大音频时长 ) print(f测试数据{len(test_dataset)} 条) # 创建数据收集器自动处理padding data_collator DataCollatorSpeechSeq2SeqWithPadding(processorprocessor) # 创建数据加载器 eval_dataloader DataLoader( test_dataset, batch_sizeargs.batch_size, num_workersargs.num_workers, collate_fndata_collator # 使用自定义收集函数 ) # 评估指标初始化 # 加载评估指标CER/字符错误率 或 WER/词错误率 metric evaluate.load(fmetrics/{args.metric}.py) # 模型评估流程 # 遍历测试数据集 for step, batch in enumerate(tqdm(eval_dataloader)): with torch.cuda.amp.autocast(): # 启用混合精度加速 with torch.no_grad(): # 禁用梯度计算评估阶段 # 音频特征处理 input_features batch[input_features].to(device_map) # 文本生成策略 if args.decoder_start_index: # 方案1使用指定数量的起始token如语言标签 print(decoder_start_index: , args.decoder_start_index) # 取labels的前N个token作为解码起始通常包含语言/任务token decoder_input_ids batch[labels][:, :args.decoder_start_index].to(device_map) # 生成文本限制最大255个新token generated_tokens model.generate( input_featuresinput_features, decoder_input_idsdecoder_input_ids, max_new_tokens255 ).cpu().numpy() else: # 方案2完全由模型自主生成起始 print(decoder_start_index: , args.decoder_start_index) generated_tokens model.generate( input_featuresinput_features, max_new_tokens255 ).cpu().numpy() # 真实标签处理 labels batch[labels].cpu().numpy() # 将-100填充标记替换为pad_token_id labels np.where(labels ! -100, labels, processor.tokenizer.pad_token_id) # 文本解码 # 将生成的token ID转为文本保留特殊token用于调试 decoded_preds processor.tokenizer.batch_decode(generated_tokens, skip_special_tokensFalse) # 将真实标签转为文本跳过特殊token decoded_labels processor.tokenizer.batch_decode(labels, skip_special_tokensTrue) # 文本后处理 # 打印原始预测和标签调试用 print(decoded_preds: , decoded_preds) print(decoded_labels: , decoded_labels) # 移除标点符号影响CER/WER计算 if args.remove_pun: decoded_preds remove_punctuation(decoded_preds) decoded_labels remove_punctuation(decoded_labels) # 繁体转简体仅中文 if args.to_simple: decoded_preds to_simple(decoded_preds) decoded_labels to_simple(decoded_labels) # 将处理后的结果添加到评估器 metric.add_batch(predictionsdecoded_preds, referencesdecoded_labels) # 内存管理 # 手动清理中间变量 del generated_tokens, labels, batch gc.collect() # 触发垃圾回收 # 计算最终结果 m metric.compute() # 计算整体错误率 print(f评估结果{args.metric.upper()}{round(m, 5)}) # 格式化输出结果 if __name__ __main__: main()