Whisperを独自データでファインチューニング - ファインチューニング編(Google Colab)

前回

こちらでデータセット（コーパス）の作成方法について説明しています。

はじめに

GoogleColabは、ファインチューニング時、38GBほどVRAMを使用しますので、お試しの際は、PROまたはPRO+を契約されておくことをおすすめします。
Whisper-Large-v3をチューニングするので結構重いです。

作ったデータセットを元に学習させる

下記のようにcorrectを埋めることができたら、これをtrain.csvとして保存して、いよいよ準備完了です。

下記のような形になるようにします。

Google Colabでファインチューニング

まずは、Google Driveと接続します。

from google.colab import drive
drive.mount('/content/drive')

次に、必要なパッケージをインストールします。

!pip install datasets>=2.6.1
!pip install accelerate -U
!pip install transformers[torch]
!pip install librosa
!pip install evaluate>=0.30
!pip install jiwer
!pip install gradio
!pip install whisper
!pip install pandas

CSVのデータをデータフレームに変換します。

import pandas as pd

# CSVファイルのパス
file_path = '/content/drive/MyDrive/ja_voice_finetuning/train.csv'

# CSVファイルを読み込む
df = pd.read_csv(file_path)

データフレームになっているか確認します。

df

下記のような感じになっているはず...

次に、学習用のデータセットと、評価用のデータセットを7:3で分けます。
評価用データセットは、validate.csvのような形で作って読み込んでもいいかもしれません。

import numpy as np
from datasets import Dataset, Audio

msk = np.random.rand(len(df)) < 0.7

train_dataset = Dataset.from_pandas(df[msk]).cast_column("path", Audio(sampling_rate=16000)).rename_column("path", "audio").remove_columns(["sampling_rate"])
validate_dataset = Dataset.from_pandas(df[~msk]).cast_column("path", Audio(sampling_rate=16000)).rename_column("path", "audio").remove_columns(["sampling_rate"])

DatasetDictにまとめて格納します。

from datasets import DatasetDict

datasets = DatasetDict({
    "train": train_dataset,
    "validate": validate_dataset
})

Processorのロード

下記で、トークナイズや特微量抽出を行うprocessorを定義します。

from transformers import WhisperProcessor
processor = WhisperProcessor.from_pretrained("openai/whisper-large-v3", language="Japanese", task="transcribe")

データセットへのmap処理のため、関数を定義します。
先ほど作ったDatasetDictをこのprocessorで前処理するprepare_dataset関数を作ります。

def prepare_dataset(batch):
    audio = batch["audio"]

    # 音響特徴量抽出
    batch["input_features"] = processor.feature_extractor(audio["array"], sampling_rate=audio["sampling_rate"]).input_features[0]

    # 正解のテキストをlabel idにエンコード
    batch["labels"] = processor.tokenizer(batch["correct"]).input_ids
    return batch

map処理をします。スペックに応じてnum_procを調整してください。
Hugging Faceブログの通りにnum_proc=4とすると途中で止まってしまうようなので、num_proc=1で実施しました。

prepared_datasets = datasets.map(prepare_dataset, remove_columns=datasets.column_names["train"], num_proc=1)

DataCollatorクラスを定義

下記のようにクラスを定義します。
詳しくは、

import torch

from dataclasses import dataclass
from typing import Any, Dict, List, Union

@dataclass
class DataCollatorSpeechSeq2SeqWithPadding:
    processor: Any

    def __call__(self, features: List[Dict[str, Union[List[int]
        , torch.Tensor]]]) -> Dict[str, torch.Tensor]:

        # 音響特徴量側をまとめる処理
        # (一応バッチ単位でパディングしているが、すべて30秒分であるはず)
        input_features \
            = [{"input_features": feature["input_features"]} for feature in features]
        batch = self.processor.feature_extractor.pad(input_features, return_tensors="pt")
        # トークン化された系列をバッチ単位でパディング
        label_features = [{"input_ids": feature["labels"]} for feature in features]
        labels_batch = self.processor.tokenizer.pad(label_features, return_tensors="pt")

        # attention_maskが0の部分は、トークンを-100に置き換えてロス計算時に無視させる
        # -100を無視するのは、PyTorchの仕様
        labels \
            = labels_batch["input_ids"].masked_fill(labels_batch.attention_mask.ne(1), -100)

        # BOSトークンがある場合は削除
        if (labels[:, 0] == self.processor.tokenizer.bos_token_id).all().cpu().item():
            labels = labels[:, 1:]

        # 整形したlabelsをバッチにまとめる
        batch["labels"] = labels

        return batch

DataCollatorを以下でインスタンス化します。

data_collator = DataCollatorSpeechSeq2SeqWithPadding(processor=processor)

評価用関数の作成

huggingfaceの公式には、日本語の評価用関数はありませんので、自分で作るしかありませんが、ありがたいことに、下記で日本語評価の方法が書いてあります。

必要なパッケージをインストール

!pip install ja-ginza
!pip install sortedcontainers~=2.1.0
!pip install spacy
import pkg_resources, imp
imp.reload(pkg_resources)

下記の、compute_metricsが評価用関数になります。

import evaluate
import spacy
import ginza

metric = evaluate.load("wer")
nlp = spacy.load("ja_ginza")
ginza.set_split_mode(nlp, "C") # CはNEologdの意らしいです

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

    # replace -100 with the pad_token_id
    label_ids[label_ids == -100] = processor.tokenizer.pad_token_id

    # we do not want to group tokens when computing the metrics
    pred_str = processor.tokenizer.batch_decode(pred_ids, skip_special_tokens=True)
    label_str = processor.tokenizer.batch_decode(label_ids, skip_special_tokens=True)

    # 分かち書きして空白区切りに変換
    pred_str = [" ".join([ str(i) for i in nlp(j) ]) for j in pred_str]
    label_str = [" ".join([ str(i) for i in nlp(j) ]) for j in label_str]

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

    return {"wer": wer}

TrainingArugumentsの定義

下記でTrainingArugumentsを定義します。
Classmethodのブログを参考に下記のように定義しました。

from transformers import Seq2SeqTrainingArguments

training_args = Seq2SeqTrainingArguments(
    output_dir="./whisper-large-v3-ja",  # change to a repo name of your choice
    per_device_train_batch_size=16,
    gradient_accumulation_steps=1,   # increase by 2x for every 2x decrease in batch size
    learning_rate=1e-5,
    # warmup_steps=500, # Hugging Faceブログではこちら
    warmup_steps=5,
    # max_steps=4000, # Hugging Faceブログではこちら
    max_steps=40,
    gradient_checkpointing=True,
    fp16=True,
    group_by_length=True,
    evaluation_strategy="steps",
    per_device_eval_batch_size=8,
    predict_with_generate=True,
    generation_max_length=225,
    # save_steps=1000, # Hugging Faceブログではこちら
    save_steps=10,
    # eval_steps=1000, # Hugging Faceブログではこちら
    eval_steps=10,
    logging_steps=25,
    report_to=["tensorboard"],
    load_best_model_at_end=True,
    metric_for_best_model="wer",
    greater_is_better=False,
    push_to_hub=False,
)

Trainerの定義

下記でtrainerを定義します。
ここで、Whisperのモデルをロードし、configを設定します。

from transformers import Seq2SeqTrainer
from transformers import WhisperForConditionalGeneration

model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-large-v3")

model.generation_config.language = "ja"


trainer = Seq2SeqTrainer(
    args=training_args,
    model=model,
    train_dataset=prepared_datasets["train"],
    eval_dataset=prepared_datasets["validate"],
    data_collator=data_collator,
    compute_metrics=compute_metrics,
    tokenizer=processor.feature_extractor,
)

学習と評価

学習前の評価を実施

# prediction_output = trainer.predict(prepared_datasets["validate"]).metrics["test_wer"]

import pandas as pd
pd.DataFrame([
    {"split":"train"
        , "wer":  trainer.predict(prepared_datasets["train"]).metrics["test_wer"]},
    {"split":"validation"
        , "wer": trainer.predict(prepared_datasets["validate"]).metrics["test_wer"]}
]
    )

この様になると思います。数字は違うかもですね。
werが低い原因としては、私がWhisperが書き起こした文章をそのまま、correctカラム（学習で使う正解ラベル）に加えているからですね。

学習

下記で学習させます。（しばらくかかります）

trainer.train()

結構負荷かかりますね

学習後の精度を確認する

下記で、学習後の精度を確認します。

import pandas as pd
pd.DataFrame([
    {"split":"train"
        , "wer":  trainer.predict(prepared_datasets["train"]).metrics["test_wer"]},
    {"split":"validation"
        , "wer": trainer.predict(prepared_datasets["validate"]).metrics["test_wer"]}
]
    )

私が、whisperが生成した内容をそのまま加えたので、あまり良い結果になっていませんが、trainのエラーレートは下がっているので、学習できているということはわかりますね。

モデルを出力

学習し終わったモデルは、下記で出力できます。あとはこれを使えばいいだけ！

trainer.save_model("/content/drive/MyDrive/ja_voice_finetuning/models")

これで、ファインチューニングは完了です。
出力したモデルで、実際に文字起こしするなどして、遊んでみるのも良いかもしれませんね！