概要

rank_bm25の初期化時にtokenizer引数を与えると、AttributeError: Can't get attribute 'tokenize' on <module '__main__' (built-in)>というエラーが出たので、解決方法を考えました。

環境

M1 Mac, macOS 14.5
Python 3.11.2
rank-bm25 0.2.2

経緯

仮想環境を作成し、rank_bm25をインストールします。

uv venv
source .venv/bin/activate
uv pip install rank-bm25

以下のスクリプトを実行します。

from rank_bm25 import BM25Okapi

corpus = [
    "The quick brown fox jumps over the lazy dog.",
    "A journey of a thousand miles begins with a single step.",
    "To be or not to be, that is the question."
]

def tokenize(text):
    return text.split()
bm25 = BM25Okapi(corpus, tokenizer=tokenize)

以下のエラーが出ます。エラーが出力され続けるのでctrl+Cで中断してください。

Process SpawnPoolWorker-23:
Traceback (most recent call last):
  File "multiprocessing/process.py", line 314, in _bootstrap
    self.run()
  File "multiprocessing/process.py", line 108, in run
    self._target(*self._args, **self._kwargs)
  File "multiprocessing/pool.py", line 114, in worker
    task = get()
           ^^^^^
  File "multiprocessing/queues.py", line 367, in get
    return _ForkingPickler.loads(res)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^
AttributeError: Can't get attribute 'tokenize' on <module '__main__' (built-in)>
Process SpawnPoolWorker-25:
Process SpawnPoolWorker-24:

原因

エラーはBM25の_tokenize_corpusメソッド内で発生しています（ソースコード）

    def _tokenize_corpus(self, corpus):
        pool = Pool(cpu_count())
        tokenized_corpus = pool.map(self.tokenizer, corpus)
        return tokenized_corpus

poolの処理に関連してエラーが発生しているようであり、軽く調べた範囲では以下の記事が原因としてもっともらしかったです。

M1 mac でmultiprocessに失敗する問題の対処法

上記記事によると、poolが標準的に使用されるメソッドがOSXの標準的なものと異なるために生じるようです（とざっくり理解しました）。
したがって、OSXの標準であるforkを使用することをPool利用時に明示的に宣言することで解決可能とのことでした。

上記を踏まえ、次のような解決策が考えられます。

解決方法

poolのcontextをforkにする

OSX限定の解決策かもしれませんが、_tokenize_corpusメソッドをオーバーライドして、poolが使用するメソッドとしてforkを指定することで、エラーが発生しなくなります。

from rank_bm25 import BM25Okapi
from multiprocessing import get_context, cpu_count
def _tokenize_corpus(self, corpus):
    with get_context("fork").Pool(cpu_count()) as pool:
        tokenized_corpus = pool.map(self.tokenizer, corpus)
    return tokenized_corpus
BM25Okapi._tokenize_corpus = _tokenize_corpus
    
corpus = [
    "The quick brown fox jumps over the lazy dog.",
    "A journey of a thousand miles begins with a single step.",
    "To be or not to be, that is the question."
]

def tokenize(text):
    return text.split()
bm25 = BM25Okapi(corpus, tokenizer=tokenize)

poolを使わない

使用メソッドを明示的に指定したくない場合は、poolの使用を避ける方法もあります。
大量の文書を処理する場合には、速度が遅くなる可能性があります。

from rank_bm25 import BM25Okapi
BM25Okapi._tokenize_corpus = lambda self, x: map(self.tokenizer, x)
    
corpus = [
    "The quick brown fox jumps over the lazy dog.",
    "A journey of a thousand miles begins with a single step.",
    "To be or not to be, that is the question."
]

def tokenize(text):
    return text.split()
bm25 = BM25Okapi(corpus, tokenizer=tokenize)

tokenizeをクラスの外で済ませる

オーバーライドをしたくない場合は、tokenizeをクラスの外で実施し、tokenizer引数は使うことを避けるという方法もあります。

from rank_bm25 import BM25Okapi

corpus = [
    "The quick brown fox jumps over the lazy dog.",
    "A journey of a thousand miles begins with a single step.",
    "To be or not to be, that is the question."
]

def tokenize(text):
    return text.split()

tokenized_corpus = [tokenize(doc) for doc in corpus]
bm25 = BM25Okapi(tokenized_corpus)

おわりに

一旦、手元で解決できる方法として３つほど検討してみました。
どれを採用するかについてですが、基本的には、3つめのクラスの外側でtokenizeする方法で良いかと思っています。
というのも、tokenizerは実は、corpusをつくるときにしか使われないため、tokenizerを初期化時に与えたところで、get_scoresなどの関数を後段で呼ぶときには、結局クラスの外側でtokenizeをする必要があるためです。
これがもし、tokenizerを初期化時に与えたらget_scoresの内部でもtokenizeを実施してくれるとかだと、1つめや2つめのオーバーライドをするのありだと思いますが。