🐯

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

2024/01/03に公開

形態素解析

LLM

jagger

tech

背景

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

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

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

作りました!

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

とりあえず使う

$ python -m pip install jagger

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

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

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

の 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)のカンマ区切りのフィールドを切り出したものです.

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

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

batch 処理(experimental)

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

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

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)でベンチマークしました.

再現スクリプトはこちらにあります.
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 の技術解説

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

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 までになります.

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

PyPI upload

cibuildwheel で python bwheel(C++ モジュール含む) を CI で一括ビルドし PyPI へアップロードするメモ

で一括でバイナリを作り(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 使っています.

セキュリティについて

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

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

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

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

WASM?

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

辞書

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

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

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

TODO

Vaporetto https://tech.legalforce.co.jp/entry/2021/09/28/180844 の Python binding https://github.com/daac-tools/python-vaporetto と速度比較をする
GPU(CUDA, OpenCL)対応して, さらなる速度の高みを目指す.
セキュリティを高める
- コア部分をセキュアに書き直す
- safetensors 形式の辞書データにする
ライセンス(利用規約)の問題が少ない辞書 or データセットから辞書を作り, Python パッケージに同梱する
J.DepP(ジョニーデップ)日本語かかり受け解析 https://zenn.dev/syoyo/articles/b4f8adeba02709 の python binding も書く(実は作業中...)

背景