微调OpenAI的Whisper自动语音识别（ASR）模型

2023/08/09

微调OpenAI的Whisper自动语音识别（ASR）模型

作者:

# Install optimum-graphcore from source 
!pip install git+https://github.com/huggingface/optimum-graphcore.git@v0.7.1 "soundfile" "librosa" "evaluate" "jiwer"
%pip install "graphcore-cloud-tools[logger] @ git+https://github.com/graphcore/graphcore-cloud-tools"
%load_ext graphcore_cloud_tools.notebook_logging.gc_logger

import os

n_ipu = int(os.getenv("NUM_AVAILABLE_IPU", 4))
executable_cache_dir = os.getenv("POPLAR_EXECUTABLE_CACHE_DIR", "/tmp/exe_cache/") + "/whisper"

# Generic imports
from dataclasses import dataclass
from typing import Any, Dict, List, Union

import evaluate
import numpy as np
import torch
from datasets import load_dataset, Audio, Dataset, DatasetDict

# IPU-specific imports
from optimum.graphcore import (
    IPUConfig, 
    IPUSeq2SeqTrainer, 
    IPUSeq2SeqTrainingArguments, 
)
from optimum.graphcore.models.whisper import WhisperProcessorTorch

# HF-related imports
from transformers import WhisperForConditionalGeneration

加载数据集

Common Voice数据集包含不同语言中演讲者朗读维基百科文本的录音。🤗 Datasets使我们能够轻松下载和准备训练并评估数据集。

首先，请确保您已经在🤗 Hub上接受了使用条款：mozilla-foundation/common_voice_13_0。一旦您接受了条款，您将拥有数据集的全部访问权限，并能够在本地下载数据。

dataset = DatasetDict()
split_dataset = Dataset.train_test_split(
    load_dataset("openslr", "SLR69", split="train", token=False), test_size=0.2, seed=0
)
dataset["train"] = split_dataset["train"]
dataset["eval"] = split_dataset["test"]
print(dataset)

需要注意的列是：

audio：原始音频样本
sentence：相应的基准转录。

我们将path列删除。

dataset = dataset.remove_columns(["path"])

由于Whisper是在16kHz采样的音频上进行预训练的，因此我们必须确保Common Voice样本相应地进行下采样。

dataset = dataset.cast_column("audio", Audio(sampling_rate=16000))

准备数据集

我们通过从原始音频输入中提取特征并注入标签来准备数据集，这些标签只是经过一些基本处理的转录。

特征提取由🤗Transformers——WhisperFeatureExtractor提供。在运行模型后将生成的标记解码为文本，我们同样需要一个标记器，WhisperTokenizer。这两者都由WhisperProcessor的实例包装。

MODEL_NAME = "openai/whisper-small"
LANGUAGE = "spanish"
TASK = "transcribe"
MAX_LENGTH = 224

processor = WhisperProcessorTorch.from_pretrained(MODEL_NAME, language=LANGUAGE, task=TASK)
processor.tokenizer.pad_token = processor.tokenizer.eos_token
processor.tokenizer.max_length = MAX_LENGTH
processor.tokenizer.set_prefix_tokens(language=LANGUAGE, task=TASK)

def prepare_dataset(batch, processor):
    inputs = processor.feature_extractor(
        raw_speech=batch["audio"]["array"],
        sampling_rate=batch["audio"]["sampling_rate"],
    )
    batch["input_features"] = inputs.input_features[0].astype(np.float16)

    transcription = batch["sentence"]
    batch["labels"] = processor.tokenizer(text=transcription).input_ids
    return batch

columns_to_remove = dataset.column_names["train"]
dataset = dataset.map(
    lambda elem: prepare_dataset(elem, processor),
    remove_columns=columns_to_remove,
    num_proc=1,
)

train_dataset = dataset["train"]
eval_dataset = dataset["eval"]

最后，我们通过使用在微调过程中将被忽略的值进行填充来预处理标签。这个填充是为了确保向模型提供具有静态形状的张量。我们通过以下数据整理器实时进行这个过程。

@dataclass
class DataCollatorSpeechSeq2SeqWithLabelProcessing:
    processor: Any

    def __call__(self, features: List[Dict[str, Union[List[int], torch.Tensor]]]) -> Dict[str, torch.Tensor]:
        batch = {}
        batch["input_features"] = torch.tensor([feature["input_features"] for feature in features])
        
        label_features = [{"input_ids": feature["labels"]} for feature in features]
        labels_batch = self.processor.tokenizer.pad(label_features, return_tensors="pt", padding="longest", pad_to_multiple_of=MAX_LENGTH)
        labels = labels_batch["input_ids"].masked_fill(labels_batch.attention_mask.ne(1), -100)

        batch["labels"] = labels

        return batch

定义指标

我们将以“错词率”（WER）来评估微调模型的性能。

metric = evaluate.load("wer")


def compute_metrics(pred, tokenizer):
    pred_ids = pred.predictions
    label_ids = pred.label_ids

    # replace -100 with the pad_token_id
    pred_ids = np.where(pred_ids != -100, pred_ids, tokenizer.pad_token_id)
    label_ids = np.where(label_ids != -100, label_ids, tokenizer.pad_token_id)

    pred_str = tokenizer.batch_decode(pred_ids, skip_special_tokens=True)
    label_str = tokenizer.batch_decode(label_ids, skip_special_tokens=True)

    normalized_pred_str = [tokenizer._normalize(pred).strip() for pred in pred_str]
    normalized_label_str = [tokenizer._normalize(label).strip() for label in label_str]

    wer = 100 * metric.compute(predictions=pred_str, references=label_str)
    normalized_wer = 100 * metric.compute(predictions=normalized_pred_str, references=normalized_label_str)

    return {"wer": wer, "normalized_wer": normalized_wer}

加载预训练模型

model = WhisperForConditionalGeneration.from_pretrained(MODEL_NAME)

model.config.max_length = MAX_LENGTH
model.generation_config.max_length = MAX_LENGTH

确保设置了适合语言的标记（如果有的话）以用于生成。我们在config和generation_config上都设置它们，以确保在生成过程中正确使用。

model.config.forced_decoder_ids = processor.tokenizer.get_decoder_prompt_ids(
    language=LANGUAGE, task=TASK
)
model.config.suppress_tokens = []
model.generation_config.forced_decoder_ids = processor.tokenizer.get_decoder_prompt_ids(
    language=LANGUAGE, task=TASK
)
model.generation_config.suppress_tokens = []

在IPU上进行Whisper的微调

可以使用IPUSeq2SeqTrainer类直接在IPU上微调模型。

IPUConfig对象指定了模型将如何在IPU之间进行流水线处理。

对于微调，我们将编码器放在两个IPU上，将解码器放在两个IPU上。

对于推理，编码器放在一个IPU上，解码器放在另一个IPU上。

replication_factor = n_ipu // 4
ipu_config = IPUConfig.from_dict(
    {
        "optimizer_state_offchip": True,
        "recompute_checkpoint_every_layer": True,
        "enable_half_partials": True,
        "executable_cache_dir": executable_cache_dir,
        "gradient_accumulation_steps": 16,
        "replication_factor": replication_factor,
        "layers_per_ipu": [5, 7, 5, 7],
        "matmul_proportion": [0.2, 0.2, 0.6, 0.6],
        "projection_serialization_factor": 5,
        "inference_replication_factor": 1,
        "inference_layers_per_ipu": [12, 12],
        "inference_parallelize_kwargs": {
            "use_cache": True,
            "use_encoder_output_buffer": True,
            "on_device_generation_steps": 16,
        }
    }
)

最后，我们指定控制训练过程的参数。

total_steps = 1000 // replication_factor
training_args = IPUSeq2SeqTrainingArguments(
    output_dir="./whisper-small-ipu-checkpoints",
    do_train=True,
    do_eval=True,
    predict_with_generate=True,
    learning_rate=1e-5 * replication_factor,
    warmup_steps=total_steps // 4,
    evaluation_strategy="steps",
    eval_steps=total_steps,
    max_steps=total_steps,
    save_strategy="steps",
    save_steps=total_steps,
    logging_steps=25,
    dataloader_num_workers=16,
    dataloader_drop_last=True,
)

然后，我们只需要将所有这些与我们的数据集一起传递给IPUSeq2SeqTrainer类：

trainer = IPUSeq2SeqTrainer(
    model=model,
    ipu_config=ipu_config,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    data_collator=DataCollatorSpeechSeq2SeqWithLabelProcessing(processor),
    compute_metrics=lambda x: compute_metrics(x, processor.tokenizer),
    tokenizer=processor.feature_extractor,
)

为了衡量WER的改进，在微调之前运行一个评估步骤。