日本語トークナイザーの作り方（トークナイザー後編）

はじめに

これまで、松尾・岩澤研究室の投稿では、何度かトークナイザーについて触れています。
前回（トークナイザー前編）は、日本語トークナイザーを準備する意味について解説しましたが、今回は実際にトークナイザーをどのように作成したか、特に

• Team JINIAC(phase1)におけるtokenizerで工夫した点
• Phase2で工夫したかった事
  について書きたいと思います。

なお、一連の手順については、トークナイザー構築のナレッジとチームの取り組み紹介【Team Kuma】ですでに詳しく述べられています。参考にされてください。

トークナイザー構築のパイプライン

データ取得について

ここでは、事前学習に使用するデータを全量使用するべきかどうか議論になりました。
確かに使用データを全量使用することで、事前学習で未知語が発生するリスクを減らすことができます。しかし、sentencepieceのbyte_fallbackオプションをTrueにすることで未知語対応はできること、低品質なデータセットに多く含まれる語彙をなるべく排除したかったことから、日本語と英語についてはwikipediaをベースにすることにしました。
一方、ヒンディー語については、そもそもwikipediaの情報量が不足していることから、学習データセットを全量使用して作成することにしました。

参考　日本語wikipediaの取得コード（google colab）

構築にあたっては、松尾・岩澤研究室から提供いただいた「標準コード」に改変を加えました。

# ドライブのマウント
from google.colab import drive
drive.mount('/content/drive')
# ベースとなるパスの指定
import os
BASE_PATH = "/content/drive/MyDrive/GENIAC/WORK/tokenizer"
%cd {BASE_PATH}

!pip install wikiextractor

import logging
import os
import bz2
import json
import shutil
import requests

# データセットダウンロード関数
def download_dataset(date: str, output_base: str = "output", lang: str = "ja") -> None:
    filename = f"{lang}wiki-{date}-pages-articles-multistream.xml.bz2"

    dump_path = os.path.join(output_base, f"tmp/wikipedia/{date}/{lang}")
    os.makedirs(dump_path, exist_ok=True)

    url = f"https://dumps.wikimedia.org/{lang}wiki/{date}/{filename}"
    if not os.path.exists(os.path.join(dump_path, filename)):
        logging.info(f"Downloading {url}")
        with requests.get(url, stream=True) as r:
            r.raise_for_status()
            logging.info(f"Saving to {os.path.join(dump_path, filename)}")
            with open(os.path.join(dump_path, filename), 'wb') as f:
                for chunk in r.iter_content(chunk_size=8192):
                    f.write(chunk)
    else:
        logging.info(f"File {os.path.join(dump_path, filename)} already exists")
        logging.info(f"Skipping download\n")

    # ダンプデータをパースする
    output_path = os.path.join(output_base, f"datasets/wikipedia/{date}/{lang}")
    if os.path.exists(output_path):
        shutil.rmtree(output_path)
    os.makedirs(output_path)

    logging.info(f"Parse and process {os.path.join(dump_path, filename)}")

# ダンプデータのダウンロード
input_date = "20240301" #最新のときは、"latest"
output_base = "dump"
lang = "ja"
status = download_dataset(input_date, output_base, lang)
# ＸＭＬファイルをテキスト形式に変換
import wikiextractor
file_name = lang + "wiki-" + input_date + "-pages-articles-multistream.xml.bz2"
dump_dir =  os.path.join(output_base, f"tmp/wikipedia/{input_date}/{lang}")
dump_path = os.path.join(dump_dir, file_name)
output_path = os.path.join(output_base, f"datasets")
!python -m wikiextractor.WikiExtractor --processes 8 -o {output_path} --no-templates {dump_path}

未知語を間引く（phase1,2での工夫）

日本語のwikipediaをベースに構築するときに、解決しておきたいことがありました。それは、ほとんど使用されないであろう単語が含まれる文章は、事前に省いておきたいということです。
では、どうやって「ほとんど使用されない」単語を特定すればよいでしょう？
今回、形態素解析器に用いられている辞書を使い、そこに登録されていない語彙についてはトークナイザー構築用のデータセットから間引くことにしました。
使用する辞書については、mecab-ipadic-NEologdやJuman++で使用されている辞書、Unidicを比較した結果、比較的最近まで更新が行われているUnidicを使用する事にしました。
また、未知語であることの判定は、fugashiを用いたときに、品詞を返してこない語としました。

参考　未知語を間引くコード（google colab）

# ドライブのマウント
from google.colab import drive
drive.mount('/content/drive')
# ベースとなるパスの指定
import os
BASE_PATH = "/content/drive/MyDrive/GENIAC/WORK/tokenizer"
%cd {BASE_PATH}
# fugashiのインストール
!pip install fugashi[unidic]
!python -m unidic download
# 未知語の間引き処理
import os
import logging
import sys
logging.basicConfig(
    stream=sys.stdout,
    level=logging.INFO,
    force=True)
def thin_out_unknown(input_path: str, output_base: str, testmode: bool) -> None:
  from fugashi import Tagger
  tagger = Tagger('-Owakati')
  input = 0
  output = 0
  flag = 0
  logging.info(f"testmode: {testmode}")
  if testmode:
    output_path = os.path.join(output_base, f"ja_wiki_thinout_test.txt")
  else:
    #output_path = os.path.join(output_base, f"ja_wiki_thinout.txt")
    output_path = os.path.join(output_base, f"ja_wiki_filter_thinout.txt")
  logging.info(f"output_path: {output_path}")
  with open(output_path, 'w', encoding='utf-8') as o:
    with open(input_path, 'r', encoding='utf-8') as f:
      for line in f:
        input += 1
        for word in tagger(line):
          if word.feature.lForm == None:
            flag += 1
        if flag == 0:
          o.write(line)
          output += 1
        flag = 0
        if input % 100000 == 0:
          logging.info(f"input: {input}, output: {output}, rate: {output / input}")
          if testmode:
            break
  logging.info(f"input: {input}, output: {output}, rate: {output / input}")
DATA_PATH = BASE_PATH + "/datasets"
text_path = DATA_PATH + "hogehoge.txt" # 使用するデータセット名に変更してください
thinout_path = os.path.join(DATA_PATH, f"thinout")
os.makedirs(thinout_path, exist_ok=True)
status = thin_out_unknown(text_path, thinout_path, testmode=False)

形態素解析を使用する（phase1のみ）

前回（トークナイザー前編）の投稿でも述べた通り、形態素解析を使用するかどうかは、メリットとデメリットが存在します。
結局、phase1では、日本語らしい区切りで学習することを優先し、形態素解析を使用しました。また、phase2では、学習速度を優先し、形態素解析は使用しませんでした。

参考　データセットに単語区切りを付与する（google colab）

Sentencepiece の分割を MeCab っぽくするを参考に作成

# ドライブのマウント
from google.colab import drive
drive.mount('/content/drive')
# ベースとなるパスの指定
import os
BASE_PATH = "/content/drive/MyDrive/GENIAC/WORK/tokenizer"
%cd {BASE_PATH}
# Mecabのインストール
!sudo apt install -y mecab libmecab-dev mecab-ipadic-utf8
# 分かち書き処理
!mecab -F"%M||||" -E"\n" input.txt -o output.txt

語彙数をいくつにするか？（事前調査）

HuggingFaceの解説では、

大きな語彙サイズは、モデルに非常に大きな埋め込み行列を入力および出力レイヤーとして持たせることを強制し、メモリおよび時間の複雑さの増加を引き起こします。一般的に、トランスフォーマーモデルは、特に単一の言語で事前トレーニングされた場合、50,000を超える語彙サイズを持つことはほとんどありません。

という記述があります。また、個人的な問題意識としては、BERT系のトークナイザーの語彙サイズは、これまで32,000である事が多く、64ビットCPUにおいて、メモリの効率的な利用を考えると、64の倍数である32,768を目安にすることは理解できます。しかし、メモリのページングを考慮すると、65,535への拡張は影響が少ないのではないかと考えました。そこで、語彙数32,000、48,000、60,000の３つのトークナイザーを作成し、事前学習速度の比較をしてみました。

実験概要

• 入力データセット　青空文庫＋法律データ＋日本語wikipedia（フィルタ済）
• 使用モデル　GPT-3 Small 125M（松尾・岩澤研究室から提供いただいた「標準コード」を改変して使用）
• 使用GPU　H100 2GPU(1node)
• 各語彙ごとに6時間処理を実施し動作確認
• 実験での確認ポイント
  • 学習速度（elapsed_ms_per_iteration,tflops）
  • tflops
  • loss
  • GPU Memory Allocated (%)

実験結果（概要）

• 学習速度は、32,000に比べ48,000と60,000は12から15%ほど遅い
  • ＜参考＞６時間で回ったiteration{ 32000 : 1410, 48000 : 1250, 60000 : 1200 }
  • ＜参考＞tfolps(max) { 32000 : 15.886, 48000 : 15.286, 60000 : 15.136 }
    
    
• loss値に顕著な違いは見られない
  • ＜参考＞６時間経過時のlm loss{ 32000 : 2.75, 48000 : 2.842, 60000 : 2.877 }
    
• メモリ確保に影響はなかった。
  

以上の実験結果から、語彙数は65,000まで増やすことを想定することにしました。

言語ごとの語彙数配分（phase1での工夫）

パイプラインにある通り、今回は言語ごとの語彙数を決めてマージをするという手法をとりました。
この方法では、ある程度自由に言語ごとの語彙数配分を決定することができます。そこで、英語・日本語・ヒンディー語の語彙数を検討することにしました。

• ＜phase1での結論＞「TOEIC満点レベルの英語を理解した大学生」レベル＋生活レベルのヒンディー語の語彙数を目指す
  • 日本語（43,000語）
    　国文学の分野では、日本人の理解語彙は、大学生で43,000～45,000語という研究があるので、これを参考にしました。（大学生の日本語の使用語彙、理解語彙）
  • 英語（13,000語）
    　多くのTOEIC対策のサイトでは、満点を取るために必要な単語数について紹介されています。その中で、最も多い単語数を提示していた、13,000語を参考にしました。
  • ヒンディー語（7,000語）
    　あまり情報は多くないのですが、「ヒンディー語の語彙本7000語」という書籍が販売されていることから、生活レベルのヒンディー語を7,000語と設定しました。

phase2で実現したかったこと

• ＜phase2での結論＞「TOEIC満点レベルの英語を理解した一般人」レベルの語彙数を目指す
  結果的に使用されなかった38Bモデル（38BモデルのLoss spikeについての反省もご覧ください。)では、phase1を改良したトークナイザーを投入していました。ヒンディー語を投入しないため、ヒンディー語用に確保していた7,000語の語彙を日本語に振り向けることで、日本語50,000語＋英語13,000語のトークナイザーになりました。
  Command R+はトークナイザーもすごかったで実験されていた、トークン数確認のプログラムで比較したところ、

トークナイザー	トークン数
LLaMA2	809
OpenAI	731
JINIAC	540
LLaMA3	516
ELYZA	474
tanuki 8×8B	450
llm-jp	434
Command R+	410
Aya	393
phase2-38B	329

と良い圧縮率を示していました。日本語に強いトークナイザーとしてお披露目したかったのですが、私自身さらにスキルアップして、次の機会を目指したいと考えています。

この成果は、NEDO（国立研究開発法人新エネルギー・産業技術総合開発機構）の助成事業「ポスト５Ｇ情報通信システム基盤強化研究開発事業」（JPNP20017）の結果得られたものです。