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

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

https://qiita.com/umaxiaotian/items/6f120919524ff8dce51b

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

# 作ったデータセットを元に学習させる
下記のようにcorrectを埋めることができたら、これをtrain.csvとして保存して、いよいよ準備完了です。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/418662/72dbd6a5-1862-cc33-6d86-8b338fc81031.png)

下記のような形になるようにします。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/418662/c1f75845-f964-0ccc-d0c2-57fec5c27e31.png)

# Google Colabでファインチューニング
まずは、Google Driveと接続します。
```python
from google.colab import drive
drive.mount('/content/drive')
```

次に、必要なパッケージをインストールします。
```python
!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のデータをデータフレームに変換します。
```python
import pandas as pd

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

# CSVファイルを読み込む
df = pd.read_csv(file_path)
```
データフレームになっているか確認します。
```python
df
```
下記のような感じになっているはず...
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/418662/4dc88722-4653-a4b2-86ee-9c019e442857.png)


次に、学習用のデータセットと、評価用のデータセットを7:3で分けます。
評価用データセットは、validate.csvのような形で作って読み込んでもいいかもしれません。
```python
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にまとめて格納します。
```python
from datasets import DatasetDict

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

## Processorのロード
下記で、トークナイズや特微量抽出を行うprocessorを定義します。
```python
from transformers import WhisperProcessor
processor = WhisperProcessor.from_pretrained("openai/whisper-large-v3", language="Japanese", task="transcribe")
```

データセットへのmap処理のため、関数を定義します。
先ほど作ったDatasetDictをこのprocessorで前処理するprepare_dataset関数を作ります。
```python
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で実施しました。

https://huggingface.co/blog/fine-tune-whisper#prepare-data

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

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

https://huggingface.co/blog/fine-tune-whisper#define-a-data-collator

```python
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を以下でインスタンス化します。
```python 
data_collator = DataCollatorSpeechSeq2SeqWithPadding(processor=processor)
```


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

https://dev.classmethod.jp/articles/whisper-fine-tuning-by-huggingface/#toc-17

必要なパッケージをインストール
```
!pip install ja-ginza
!pip install sortedcontainers~=2.1.0
!pip install spacy
import pkg_resources, imp
imp.reload(pkg_resources)
```

下記の、compute_metricsが評価用関数になります。
```python
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のブログを参考に下記のように定義しました。
```python
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を設定します。
```python
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,
)

```


## 学習と評価
### 学習前の評価を実施

```python
# 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カラム（学習で使う正解ラベル）に加えているからですね。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/418662/7bc2d4e2-e179-d56c-eeb0-85485b6d69d8.png)


### 学習
下記で学習させます。（しばらくかかります）
```python
trainer.train()
```
結構負荷かかりますね
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/418662/74edd1b5-2066-b375-84c6-238d0d826440.png)

### 学習後の精度を確認する
下記で、学習後の精度を確認します。
```python
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のエラーレートは下がっているので、学習できているということはわかりますね。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/418662/47ca58df-ae28-7988-efde-26120484c5d5.png)

### モデルを出力
学習し終わったモデルは、下記で出力できます。あとはこれを使えばいいだけ！
```python
trainer.save_model("/content/drive/MyDrive/ja_voice_finetuning/models")
```


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