pickle.dump できる natto.MeCab のラッパーを作ってみた

6 min read読了の目安(約6000字

この記事は freee データに関わる人たち Advent Calendar 2020 / 4 日目のエントリーです。中 1 日での登板です。

はじめに

2 日目の記事では

私は 4 日目にも登場する予定です。4 日目は徹夜せずに書き終えたいですw

と書きましたが結局書き始めたのは夜中でした 😇 自分の無計画さを呪いたくなります。

さて今回は MeCab Python バインディングの一実装である natto-pypickle.dump() できるようにした超簡易ラッパー ShioKobu を作ったのでその紹介をしようと思います。(作ったのは結構前です)

簡易すぎてチビるで?

natto.MeCab は picklable ではない

natto-py は MeCab の Python バインディング[1]です。

Python から MeCab を使うときにとても便利なライブラリなのですがこの natto-py における natto.MeCab[2] はそのままでは pickle.dump() できません。

In [1]: import natto

In [2]: import pickle

In [2]: args = dict(output_format_type="wakati")

In [3]: tagger = natto.MeCab(args)

In [4]: tagger.parse("すもももももも")
Out[4]: 'すもも も もも も'

In [5]: _ = pickle.dumps(tagger)
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-5-ba94ac1189e4> in <module>
----> 1 _ = pickle.dumps(tagger)

TypeError: cannot pickle 'module' object

それだけ聞くとふーんと思われるかも知れませんが。Python で MeCab を使った NLP 関連の処理を書いていると natto.MeCabMeCab.Tagger のオブジェクトをインスタンス変数として持っているオブジェクトを pickle.dump() したくなる場面に遭遇します。

例えば以下の Tokenizer クラスのようにストップワードを持たせてトークナイズ結果をフィルタリングするコードがあったとします。

In [1]: import natto

In [2]: class Tokenizer(object):
    ...:     def __init__(self, stopwords: set = set()):
    ...:         self.tagger = natto.MeCab()
    ...:         self.stopwords = stopwords
    ...: 
    ...:     def __call__(self, text):
    ...:         ret = []
    ...:         for node in self.tagger.parse(text, as_nodes=True):
    ...:             if node.is_eos():
    ...:                 return ret
    ...:             if (
    ...:                 node.feature.startswith("名詞")
    ...:                 and node.surface not in self.stopwords
    ...:             ):
    ...:                 ret.append(node.surface)
    ...: 

In [3]: tokenize = Tokenizer(stopwords={"もも"})

In [4]: tokenize("すもももももも")
Out[4]: ['すもも']

で, トークナイズしたい文書がたくさんあるから joblib を使って並列実行しようとか思って以下のようなコードを書くと joblib は tokenizenatto.MeCab をインスタンス変数に持った Callable なオブジェクト) を pickle.dump() しようとするのでエラーになるわけです。

In [5]: from joblib import delayed, Parallel

In [6]: docs = ["すもももももも", "なのは完売", ...]  # でかい文書集合 

In [7]: Parallel(n_jobs=2)(delayed(tokenize)(doc) for doc in docs)
---------------------------------------------------------------------------
...
PicklingError: Could not pickle the task to send it to the workers.

回避方法は色々思いつくけどなるべくコードに変更を入れずに picklable にしたい。そんな用途に応えるため作ったのが ShioKobu です。

ShioKobu で MeCab を picklable にする

ShioKobu[3] は natto-py の超超薄いラッパーです。 仕組みは簡単で natto.MeCab を継承して __reduce_ex__() を生やしてやっただけです。__reduce_ex__() メソッドについては以下の記事を参考にしました。

__reduce_ex__() を定義してやると pickle.dump(), pickle.load() したときの挙動をある程度制御することができます。参考にした記事にも書かれていますが

  • pickle.dump() したとき
    • __reduce_ex__() が呼ばれる
  • pickle.load() したとき
    • __reduce_ex__() の戻り値が __init__() に渡されてオブジェクトが作られる

のような挙動になるので natto.MeCab を直接 pickle.dump() しなくてすむようになります。また, 工夫というほどの工夫ではないですが ShioKobu では

  1. __init__() の引数を NamedTuple に詰めて self.__args へ退避させる, 後段で分岐させやすくするために NamedTuple で独自型にしてやる
  2. __reduce_ex__()self.__argspickle.load() 時の引数として返す
  3. __reduce_ex__() の戻り値は __init__() の第一引数に渡ってくるのでそこを型チェックして pickle 経由で呼び出された場合と普通の初期化で呼び出された場合の処理を分岐させる

のように実装して, 元々あった natto.MeCab の初期化インタフェースをそのまま使えるようにしました。

https://github.com/nagomiso/shiokobu/blob/master/shiokobu/mecab.py

使い方は超簡単。 natto の代わりに shiokobu を import するだけです。

In [1]: import shiokobu

In [2]: class Tokenizer(object):
    ...:     def __init__(self, stopwords: set = set()):
    ...:         self.tagger = shiokobu.MeCab()
    ...:         self.stopwords = stopwords
    ...: 
    ...:     def __call__(self, text):
    ...:         ret = []
    ...:         for node in self.tagger.parse(text, as_nodes=True):
    ...:             if node.is_eos():
    ...:                 return ret
    ...:             if (
    ...:                 node.feature.startswith("名詞")
    ...:                 and node.surface not in self.stopwords
    ...:             ):
    ...:                 ret.append(node.surface)
    ...: 

In [3]: tokenize = Tokenizer(stopwords={"もも"})

In [4]: tokenize("すもももももも")
Out[4]: ['すもも']

In [5]: from joblib import delayed, Parallel

In [6]: docs = ["すもももももも", "なのは完売", ...]

In [7]: Parallel(n_jobs=2)(delayed(tokenize)(doc) for doc in docs)
Out[7]: [['すもも'], ['の', '完売']]

修正部分も少なく, かつエラーが発生せずに無事に動いていますね。

※ ただし, これは natto.MeCab がシリアライズされているわけではなく各プロセス上で新たにオブジェクトが生成されていることに注意してください。natto.MeCab のような初期化コストが大きめのオブジェクトがボコボコ作られることになるのでオーバーヘッド込みで考えると並列化しても早くならない可能性があります。

終わりに

というわけで natto-py を使う人で最小の変更で picklable にするための超薄いラッパー ShioKobu を作ってみました。使う人は自分ぐらいしかいない気がしますが __reduce_ex__() を定義して普通は pickle できないオブジェクトを pickle できるようにするためのサンプルコードぐらいにはなるかなと思います。

ちなみにこの記事を書いている途中で 【Python】MeCabのTaggerオブジェクトを持つ単語分割器をpickleで保存する方法
という記事を見つけて 『内容もろ被っとるやん 😇 』 となったことをここに記します。

5 日目は harry さんの 『アプリエンジニアチームからデータエンジニアチームに異動した話(仮)』 です。お楽しみに。

脚注
  1. Python バインディングとしは mecab-python3 の方が有名かと思いますが何となく natto-py のほうが API が好みです。 ↩︎

  2. mecab-python3 だと MeCab.Tagger↩︎

  3. MeCab からの海藻つながりで塩漬けにしたら塩昆布だなと安直につけた名前です。 ↩︎