🐯

高速形態素解析 Jagger の Python binding のメモ

2024/01/03に公開

背景

LLM 用データセット(コーパス)構築で, 多量のテキストデータ(10 TB くらい)を形態素解析する必要がある... 既存のでは遅すぎ...(GiNZA, Sudachi). すごい精度はいらないが速いのがほしい...

C++ での史上最速(たぶん)で形態素解析できる jagger

https://zenn.dev/syoyo/articles/8381ddef921a5a

あるけど, python binding ほしいぽよ...

作りました!

https://github.com/lighttransport/jagger-python

(C++ での)マルチスレッドでバッチ処理するのも作って, 速度の高みを更に目指しました!

とりあえず使う

$ python -m pip install jagger

でインストールいけます!

Windows(ARM も!), Linux(arm も!), macOS の全バイナリがあるので, コンパイル不要でいけるよ.

辞書データは別途必要です. とりあえずは

https://github.com/lighttransport/jagger-python

の README からたどれるコンパイル済み KWDLC を使いましょう.

import jagger

model_path = "model/kwdlc/patterns"

tokenizer = jagger.Jagger()
tokenizer.load_model(model_path)

text = "吾輩は猫である。名前はまだない。"
toks = tokenizer.tokenize(text)

for tok in toks:
    print(tok.surface(), tok.feature())
print("EOL")
吾輩    名詞,普通名詞,*,*,吾輩,わがはい,代表表記:我が輩/わがはい カテゴリ:人
は      助詞,副助詞,*,*,は,は,*
猫      名詞,普通名詞,*,*,猫,ねこ,*
である  判定詞,*,判定詞,デアル列基本形,だ,である,*
。      特殊,句点,*,*,。,。,*
名前    名詞,普通名詞,*,*,名前,なまえ,*
は      助詞,副助詞,*,*,は,は,*
まだ    副詞,*,*,*,まだ,まだ,*
ない    形容詞,*,イ形容詞アウオ段,基本形,ない,ない,*
。      特殊,句点,*,*,。,。,*

Voila~

# print tags
for tok in toks:
    # print tag(split feature() by comma)
    print(tok.surface())
    for i in range(tok.n_tags()):
        print("  tag[{}] = {}".format(i, tok.tag(i)))

で tag(feature のフィールド)を取得できます.

tag

feature 情報(CSV)のカンマ区切りのフィールドを切り出したものです.

https://hayashibe.jp/tr/mecab/dictionary/

にあるように, 辞書によってはフィードがクオートされている場合もあるので, そのあたりも考慮しています(quote 文字を指定可能. デフォルトは ")

速さのさらなる高みへ...

batch 処理(experimental)

複数行をマルチスレッドで一括処理する tokenize_batch を実装しました.

入力行をマルチスレッドで改行で分解するテクニックは,
CPU での史上最速の CSV parser のひとつ nanocsv のを利用しました.
(GPU で CUDA で CSV parse するのがあるので, CPU に限定しました 😌)

https://github.com/lighttransport/nanocsv

feature 文字列(CSV)の分割にも nanocsv のコードの一部を使っています.

import jagger

model_path = "model/kwdlc/patterns"

tokenizer = jagger.Jagger()
tokenizer.load_model(model_path)

text = """
吾輩は猫である。
名前はまだない。
明日の天気は晴れです。
"""

# optional: set C++ threads(CPU cores) to use
# default: Use all CPU cores.
# tokenizer.set_threads(4)

toks_list = tokenizer.tokenize_batch(text)

for toks in toks_list:
    for tok in toks:
        print(tok.surface(), tok.feature())

でいけます. たぶん形態素解析よりも, 処理結果の文字列を Python world にもってくるほうが時間かかるかも...

メモリはそれなりに使います(入力文字列の 60 倍くらい). したがって一括で処理よりは, 10 万行ずつとかが推奨です. たとえば 1024*128 で 12 万行でおよそ 8 GB ほど消費しました.

benchmark

日本語 wikipedia 40b 371 万行(1.2 GB)でベンチマークしました.

https://github.com/lighttransport/jagger-python/tree/main/benchmark

再現スクリプトはこちらにあります.
Ryzen 3950X(16 cores) で測定しました.

import jagger
import tqdm
import time

model_path = "model/kwdlc/patterns"

tokenizer = jagger.Jagger()
tokenizer.load_model(model_path)

lines = open("output-wiki.txt", 'r', encoding='utf8').readlines()

s = time.time()
for line in tqdm.tqdm(lines):
    toks = tokenizer.tokenize(line)

e = time.time()
print("Jagger: Total {} secs".format(e - s))

Jagger: Total 203.34391069412231 secs

ふうむ... オフィシャルサイト https://www.tkl.iis.u-tokyo.ac.jp/~ynaga/jagger/index.en.html では M2 macbook で 1 core で 100 万行を 1 秒で処理できるとありますから, 結構な速度低下です...
top で見てもあまり CPU cores を使っていないので, C++ <-> Python でのデータ変換が結構時間を食っていると思われます.

ちなみに batch_tokenize のほうが速度が遅い結果となりました...
(C++ <-> Python 変換(シングルスレッド)に時間がかかっているようです)

マルチプロセッシング

C++ コア部分ではロック関連の問題は無いので, 速さの高みを目指すなら Python でマルチプロセッシングで tokenize するのがよいでしょう.

以下のような感じです. ここでは tokenize_bach を使うようにしてみました.

だいたい Wikipedia 1.2 GB(371 万行) が Ryzen 3900X 12 コアマシン(行単位処理のベンチのときとは別 PC)で 30 秒で処理できました(タグ文字列取得など除く). これだと実用的で, 1 TB テキストを 1 Ryzen で 1 日で形態素解析できそうですかね.

ちなみに, 現状では tokenize, tokenize_batch でかえって来る Tag クラスの Picking 対応していないため, 処理結果をそのまま future の結果として返そうとするとエラーになります.
(cuncurrent futures では future の結果を picking して返すらしい)

したがって処理関数の中でなにやらうまく処理したり, future の結果としてなにかほしかったら pure Python なオブジェクトを作って返すようにしてください.

  import jagger
  import concurrent.futures
  from multiprocessing import cpu_count
  import os
  import sys
  from tqdm import tqdm

  model_path = "model/kwdlc/patterns"

  tokenizer = jagger.Jagger()
  tokenizer.load_model(model_path)
  #tokenizer.set_threads(16)

  lines = open("output-wiki.txt", 'r', encoding='utf8').readlines()

  # Use half of CPU cores
  num_process = max(1, cpu_count() // 2)

  nlines_per_batch = 1000

  def run(lines):
      # TODO: Accept List[str] as input for tokenize_batch
      toks_list = tokenizer.tokenize_batch(''.join(lines))

      # NOTE: Cannot return tokenized result at the moment. List[List[PyToken]]] fails pickle serialization
      # So process toks_list here and convert to pure Python object if you want to return something.
      return None
   
  total_ticks = max(1, len(lines) // nlines_per_batch)
  with tqdm(total=total_ticks) as pbar:
      with concurrent.futures.ProcessPoolExecutor(max_workers=num_process) as executor:
          futures = {executor.submit(run, lines[i:i+nlines_per_batch]): i for i in range(0, len(lines), nlines_per_batch)}

          results = {}
          for future in concurrent.futures.as_completed(futures):
              arg = futures[future]
              result = future.result()
              pbar.update(1)

Vaporetto との比較

速度の高みを目指す:高速な単語分割器 Vaporetto の技術解説
https://tech.legalforce.co.jp/entry/2021/09/28/180844

これはいい対戦相手ですね.

import time

import vaporetto
import zstandard
import tqdm

dctx = zstandard.ZstdDecompressor()
with open('bccwj-suw+unidic_pos+pron/bccwj-suw+unidic_pos+pron.model.zst', 'rb') as fp:
    with dctx.stream_reader(fp) as dict_reader:
        tokenizer = vaporetto.Vaporetto(dict_reader.read(), predict_tags = True)

lines = open("output-wiki.txt", 'r', encoding='utf8').readlines()

s = time.time()
for line in tqdm.tqdm(lines):
    toks = tokenizer.tokenize(line)

e = time.time()

print("Vaporetto: Total {} secs".format(e - s))
Vaporetto: Total 190.48942708969116 secs

なんと! 行単位で処理の場合 Vaporetto のほうが速い結果となりました!

行単位だと Python 界のオーバヘッドが大きそう + 辞書が違うというのもありますが, いずれにせよjagger-python の C++ <-> Python の変換を見直したほうがよさそうですね.

tag 情報などの文字列返しの高速化(TODO)

surface, feature, tag は, 現状 Python 界へ文字列として返すために毎回 std::string を作っています.
(jagger C++ では, feature の文字列は辞書データへのポインタで対応している)

C++17 string_view で効率化ができますが, pypi 用の manylinux がデフォですと manylinux2014 となり, CentOS7 対応のために結構古い glibc 使っていたりで, よくて C++14 までになります.

https://qiita.com/syoyo/items/d7d1cc6e8ad48326d792

一応 const char* を Python 界に byte 列で返すみたいなやり方があった気がするので, それを使うか, 品詞などは enum にするなどがよいでしょう
(Python で tag == "動詞" とか書くのめんどいですしね)

PyPI upload

cibuildwheel で python bwheel(C++ モジュール含む) を CI で一括ビルドし PyPI へアップロードするメモ
https://qiita.com/syoyo/items/97f35b4d5c40761cc314

で一括でバイナリを作り(Linux aarch64, Windows ARM もあるよ), PyPI upload しています.

最近の PyPI upload は, twine で API token でアップロードではなく, Trusted Publisher が推奨になったようです. 現時点(2023/01)では Github Actions のみの対応のようなので, Github Actions でのみ Trusted Publisher でのアップロードにしています.
(詳細は git repo の .github/workflows/wheel.yml 参照してね)

linux arm 版は, Github Actions だと qemu ビルドで遅いので, Cirrus CI 使っています.

https://zenn.dev/syoyo/articles/2ee2451bdc3d3f

セキュリティについて

我々が書いた部分はそこそこ考慮していますが, 元の jagger のコア部分は未対応です.

したがって変な辞書データ(malcious input)を与えられたらクラッシュするでしょう.

そのうちコア部分も書き直してセキュアにしたいですね.
(コア部分はそんなに分量ないですが, めちゃ読みづらいコードなので読解に時間がかかる...)

辞書データのフォーマットも, safetensors 形式にして扱いやすいようにしたいところです.

https://zenn.dev/syoyo/articles/50473523cfa376

WASM?

pybind11 のコードを流用して, WASM 版も比較的用意に作れるでしょう.
ブラウザで高速形態素解析が可能になります!

辞書

KWDLC はライセンス(利用規約)がありません.

https://github.com/ku-nlp/KWDLC/issues/37

したがって利用は自己責任になります.

Python パッケージに辞書も同梱したいのを考えると, UniDic(mecab 形式)を変換するか, wikipedia dump などを GiNZA electra で形態素解析してコーパスを作り辞書を作るのがよいでしょう.

TODO

Discussion