自動タグ付け機能でも作ってみる - 7: BM25の実装

前提

先の構成図 参照。

本丸、 KeywordRanker を実装していく。

Okapi BM25を分解する

Okapi BM25 による、文書d、単語wの重み付け (Combined Weight: CW) は、次のような式で表される。

$$ CW(w,d) = IDF(w) \cdot TF(w,d) \cdot \frac{k_1 + 1}{k_1 \cdot (1 - b + (b \cdot NDL(d))) + TF(w,d)} $$

まず、チューニングパラメータが2つある。

  • k1: 単語出現頻度に係るチューニングパラメータ。
    • $ k_1 \in [1.2,2.0] $
    • 今回は 2.0 とする
  • b: 文書の長さに係るチューニングパラメータ。
    • $ 0 < b < 1 $
    • 今回は 0.75 とする

また、数式中の各関数は、次のような意味を持つ。

  • IDF (Inverse Document Frequency): 単語の出現する文書の頻度。
    • $ IDF(w) = log( \frac{\text{文書の総数}}{\text{単語wが出現する文書数}} ) $
  • TF (Term Frequency): 文書中における出現頻度の高さ。
    • $ TF(w, d) = \frac{\text{単語wの文書dにおける出現回数}}{\text{文書dの単語数}} $
  • NDL (Normalized Document Length): ある文書の単語数の文書群全体における多さ。
    • $ NDL(d) = \frac{\text{文書dの単語数}}{\text{すべての文書の平均単語数}} = \frac{\text{文書dの単語数} \cdot \text{文書の総数}}{\text{総単語数}} $

つまり、 KeywordRanker (WordStorage) は、文書と単語ごとに次の5つを集計していけばいい。

  • 単語wの文書dにおける出現回数
  • 文書dの単語数
  • 総単語数
  • 文書の総数
  • 単語wが出現する文書数

単語集計の実装

後者2つは他の集計するべき項目に従属しているので、簡略化すると…

from math import log
from typing import Dict, List
from collections import defaultdict as dd

from .word import Word
from .word_storage import WordStorage


class KeywordRanker(WordStorage):

    def __init__(self) -> None:
        super().__init__()

        # 単語wの文書dにおける出現回数
        self.query_freq: Dict[Word, Dict[str, int]] = dd(lambda: dd(int))
        # 文書dの単語数
        self.doc_length: Dict[str, int] = dd(int)
        # 総単語数
        self.word_count: int = 0
        # 文書の総数 = len(self.doc_length)
        # 単語wが出現する文書数 = len(self.query_freq[w])


    def push(self, path: str, word: Word) -> None:
        self.doc_length[path] += 1
        self.query_freq[word][path] += 1
        self.word_count += 1

となる。

BM25 Combined Weightの実装

集計された値を元に、すべての文書、すべての単語ごとの重みは次のように得られる。

def weighted_keywords(self) -> Dict[str, List[Word]]:
    K1 = 2
    B = 0.75

    # 文書の総数
    doc_count = len(self.doc_length)

    '''
    すべての文書の平均単語数の逆数
    (文書の総数) / (総単語数)
    '''
    adli = doc_count / self.word_count

    ret: Dict[str, List[Word]] = dd(lambda: [])

    for word, freq in self.query_freq.items():
        idf = log(doc_count / len(freq))

        for path, count in freq.items():
            tf = freq[path] / self.doc_length[path]
            ndl = self.doc_length[path] * adli

            props = word._asdict()
            props["weight"] = (tf * idf * (K1 + 1)) / \
                (K1 * (1 - B + (B * ndl)) + tf)

            ret[path].append(Word(**props))

    for path, words in ret.items():
        ret[path] = sorted(words, key=lambda w: -w.weight)

    return ret

ただし、重みを単語 (Word) に突っ込んでいるので、 Word クラスの定義に weight を追加している。

class Word(namedtuple('Word', (
    'word',                # 原形
    'feature_id',          # 素性ID
    'weight',              # 重み付け
))):
    pass

実際動かしてみると、2時間ほどかかった上で、大量の単語辞書が得られた。 出力が多すぎたので、割愛。

残課題

Wikipediaの集計に2時間かかるので、

  • 集計結果をしかるべき形でDump
  • タグとなるキーワードの抽出を行いたい文書を投入する際に、Restore

する仕組みを作る必要がある。