HuggingSoundによる音声認識モデルのfine-tuning

こんにちは
AIチームの戸田です

今回は最近リリースされたHuggingFaceをベースにした音声処理用のツール、HuggingSoundを使って音声認識を試してみます。

HuggingSoundはHuggingFaceが公式に出しているものではなく、リオデジャネイロ大学のJonatas Grosman氏が個人的に開発しているライブラリで、今年に入ってリリースされたばかりの新しいライブラリです。

日本語の音声認識モデルはGrosman氏がこのライブラリを使ってCommon Voice、 CSS10、 JSUTをfine-tuningしたものを公開してくれていますが、本記事ではイチからfine-tuningを試してみたいと思います

データ準備

学習に使用するデータは声庭で公開されているデータを利用します。声庭は有志でアノテーションを行っている公開音声データで、利用・修正・再配布が自由なオープンデータになっています。

2022年4月現在にアノテーションされているデータは、まだ1000フレーズ程度なので、End-to-Endな日本語認識ではなく、日本語の音素認識を試してみたいと思います。

HuggingSoundで学習するには、以下のように音声ファイルのPATHとその音声ファイルの書き起こしの辞書のリストを入力する必要があります。

references = [
    {"path": "/path/to/sagan.mp3", "transcription": "extraordinary claims require extraordinary evidence"},
    {"path": "/path/to/asimov.wav", "transcription": "violence is the last refuge of the incompetent"},
]

以下に声庭のデータからHuggingSoundで学習するためのデータセットの形式に変換するコードを記載します。音素抽出にはpyopenjtalkを利用します。

import json
import random
import librosa
import pyopenjtalk
import pandas as pd
import soundfile as sf

# グローバル設定
class GCF:
    # git cloneで取得した声庭ファイル
    KONIWA_ROOT = './koniwa/data'
    # アノテーション済みのデータリスト
    ANNOTATION_DATA = [
        'amagasaki/amagasaki__2011_04_20',
        'free_culture_2012/free_culture_2012__1',
        'tnc/tnc__gongitsune',
        'tnc/tnc__kumonoito',
        'tnc/tnc__nezuminoyomeiri',
        'tnc/tnc__tebukurowokaini',
    ]
    # モデルの出力先
    OUTPUT_DIR = f"./output/finetuned"
    # 前処理済みデータの保存先
    PREPROCESSED = f"./output/preprocessed"
    # 評価データ数
    N_EVAL = 200

audio_data = []
for annotation in GCF.ANNOTATION_DATA:
    d = json.load(open(f'{GCF.KONIWA_ROOT}/{annotation}.json', 'r'))
    df = pd.DataFrame(d['annotation'])
    # pyopenjtalkを使って音素変換
    df['phoneme'] = df['data'].map(lambda x: pyopenjtalk.g2p(x['kana_level0'], kana=False).lower())
    # tncの音声データのみ拡張子が.aacでそれ以外は.mp3
    suffix = 'aac' if 'tnc' == annotation[:3] else 'mp3'
    a, sr = librosa.load(f"./{annotation}.{suffix}")
    # フレーズごとに音声ファイルを分割する
    for idx, row in df.iterrows():
        i = int(row.start * sr)
        j = int(row.end * sr)
        speech = a[i:j]
        fname = annotation.split('/')[-1]
        p = f'{GCF.PREPROCESSED}/{fname}_{idx:04}.wav'
        sf.write(p, speech, sr, 'PCM_24')
        r = {"path": p, "transcription": row['phoneme']}
        audio_data.append(r)

df = pd.DataFrame(audio_data).dropna()
audio_data = [row.to_dict() for _, row in df.iterrows()]
# 短すぎる音声は除外する
audio_data = [i for i in audio_data if len(i['transcription']) > 10]

# 学習データと評価データに分割
random.shuffle(audio_data)
eval_data, train_data= audio_data[:GCF.N_EVAL], audio_data[GCF.N_EVAL:]

これでデータ準備は完了です。

学習

公開されているwav2vec2-largeをfine-tuningします。パラメータは決め打ちです。

model = SpeechRecognitionModel("facebook/wav2vec2-large-xlsr-53", device=device)
tokens = list('abcdefghijklmnopqrstuvwxyz ')
token_set = TokenSet(tokens)

training_args = TrainingArguments(
    learning_rate=3e-4,
    max_steps=1000,
    eval_steps=200,
    per_device_train_batch_size=2,
    per_device_eval_batch_size=2,
)
model_args = ModelArguments(
    activation_dropout=0.1,
    hidden_dropout=0.1,
)

model.finetune(
    GCF.OUTPUT_DIR, 
    train_data=train_data, 
    eval_data=eval_data,
    token_set=token_set, 
    training_args=training_args,
    model_args=model_args,
)

こちらのコードをGoogle ColabのTesla K80で実行しました。学習時間はおおよそ30分くらいでした。

HuggingSoundでは評価も簡単に行なえます。以下のコードで評価データの文字誤り率(CER)と単語誤り率(WER)を計算することができます。

evaluation = model.evaluate(eval_data, inference_batch_size=1)
print(evaluation)
# {'cer': 0.22771011877913097, 'wer': 0.282102666872206}

両方とも2割以上間違えているので、あまり良い精度とは言えませんね。

最後にモデルを使った推論をしてみます。私の声で録音した「株式会社AI Shift(かぶしきがいしゃえーあいしふと)」という音声に対して今回学習したモデルで音素を予測してみます。

録音した音声(test_audio.mp3)
audio_paths = ['test_audio.mp3']
transcriptions = model.transcribe(audio_paths, batch_size=1)

認識結果は以下のようになりました。

'k k a b u sh i k i r a e s a a i sh i s u t o'

私の発音が悪いのもあるかもしれませんが、会社(がいしゃ)が「らえさ(ra e s a)」になっていたりと、認識ミスはあるものの、ざっくりとは音素認識できているように思えます。

1000フレーズ程度でこれだけの精度が出せるので、データを増やしたり、パラメータをチューニングすればもっと精度はよくなると思います。

おわりに

本記事ではHuggingSoundを使った音声認識を試してみました。

以前HuggingFaceでwav2vecのfine-tuningを行ったことがあるのですが、DataLoader周りの実装が少し面倒だったことを覚えています。今回のHuggingSoundだと音声ファイルのパスを指定するだけで良いので、非常に簡単にfine-tuningが実行できました。

まだリリースされたばかりのツールなので、ドキュメントやサンプルコードなどが整備されていないようですが、音声認識に関してはデータさえあれば簡単に学習できそうです。音声認識以外も増やしていくようなので、今後のバージョンアップに期待ですね。

最後までお読みいただきありがとうございました!