Tutorial (Japanese BM25 and Document Search)

Implement BM25 for Japanese text with nagisa

This tutorial provides an example of implementing BM25 for Japanese text by using nagisa for tokenization and nagisa.stopwords for removing stopwords. You extract keywords from news articles with BM25, and build a simple Japanese document search engine. Finally, you use nagisa with BM25S (BM25-Sparse), a fast BM25 library.

Install python libraries

Before we get started, please run the following commands to install the libraries used in this tutorial.

pip install nagisa
pip install tqdm
pip install bm25s

Note that bm25s requires Python 3.8 or later.

How document search works

A document search engine returns the documents relevant to a query in order of relevance. It splits documents and queries into words (terms), builds an inverted index that maps each term to the documents containing it, computes a score for each document that contains the query terms, and returns the top-k documents with the highest scores. Japanese text has no spaces between words, so the first step needs a tokenizer such as nagisa. BM25 is a standard function for computing this score.

BM25

BM25 gives a high score to a document that contains the query terms many times, especially the terms that are rare in the corpus. The score of a document \(d\) for a query \(q\) is

\[\mathrm{score}(q, d) = \sum_{t \in q} \mathrm{IDF}(t) \cdot \frac{\mathrm{tf}(t, d)}{\mathrm{tf}(t, d) + k_1 \left(1 - b + b \frac{|d|}{\mathrm{avgdl}}\right)}\]

\[\mathrm{IDF}(t) = \ln \left(1 + \frac{N - \mathrm{df}(t) + 0.5}{\mathrm{df}(t) + 0.5}\right)\]

where \(\mathrm{tf}(t, d)\) is the frequency of the term \(t\) in the document \(d\), \(\mathrm{df}(t)\) is the number of documents containing \(t\), \(N\) is the number of documents, \(|d|\) is the length (the number of terms) of \(d\) and \(\mathrm{avgdl}\) is the average document length.

• \(k_1\) (usually 1.2-2.0) controls how quickly the score saturates as the term frequency increases.

• \(b\) (usually 0.75) controls the normalization by the document length.

This is the variant used by Lucene and BM25S. The original Okapi BM25 also multiplies each summand by \((k_1 + 1)\), which does not change the ranking, and its IDF can become negative.

Download the dataset

This tutorial uses the Livedoor News Corpus, which contains 7,367 Japanese news articles in 9 categories. The corpus (~9 MB) is downloaded automatically into the ldcc/ directory on the first run. Each document is the title and the body of an article.

Tokenize Japanese text with nagisa

split_words() keeps the content words (nouns, verbs, adjectives, etc.) using nagisa.extract(). nagisa works on a sentence, so the text is tokenized line by line. The words are lowercased so that a query “iphone” matches “iPhone”. tokenize() then removes the stopwords in nagisa.stopwords.

def split_words(text):
    # nagisa works on a sentence, so tokenize the text line by line.
    words = []
    for line in text.splitlines():
        tokens = nagisa.extract(line, extract_postags=TARGET_POSTAGS)
        words += [word.lower() for word in tokens.words]
    return words

def tokenize(text):
    return [word for word in split_words(text) if word not in nagisa.stopwords]

Particles, auxiliary verbs and symbols are removed based on their POS-tags, and “し” is removed as a stopword.

print(nagisa.tagging("iPhoneで使える便利なアプリを紹介します。"))
#=> iPhone/名詞 で/助詞 使える/動詞 便利/名詞 な/助動詞 アプリ/名詞 を/助詞 紹介/名詞 し/動詞 ます/助動詞 。/補助記号

print(tokenize("iPhoneで使える便利なアプリを紹介します。"))
#=> ['iphone', '使える', '便利', 'アプリ', '紹介']

Implement BM25

The BM25 class builds an inverted index (term -> {doc_id: term frequency}) and the IDF of each term in fit(). term_weight() computes the BM25 weight of a term in a document, which is the summand of the formula above.

class BM25:

    def __init__(self, k1=1.5, b=0.75):
        self.k1 = k1
        self.b = b

    def fit(self, corpus_tokens):
        self.num_docs = len(corpus_tokens)
        self.doc_lens = [len(tokens) for tokens in corpus_tokens]
        self.avgdl = sum(self.doc_lens) / self.num_docs

        # Inverted index: term -> {doc_id: term frequency}
        self.index = {}
        for doc_id, tokens in enumerate(corpus_tokens):
            for term, tf in Counter(tokens).items():
                self.index.setdefault(term, {})[doc_id] = tf

        # Inverse document frequency of each term
        self.idf = {}
        for term, postings in self.index.items():
            df = len(postings)
            self.idf[term] = math.log(1 + (self.num_docs - df + 0.5) / (df + 0.5))
        return self

    def term_weight(self, term, doc_id):
        tf = self.index.get(term, {}).get(doc_id, 0)
        if tf == 0:
            return 0.0
        norm = self.k1 * (1 - self.b + self.b * self.doc_lens[doc_id] / self.avgdl)
        return self.idf[term] * tf / (tf + norm)

    def search(self, query_tokens, k=5):
        # Only documents that contain a query term get a score.
        scores = Counter()
        for term in query_tokens:
            for doc_id in self.index.get(term, {}):
                scores[doc_id] += self.term_weight(term, doc_id)
        return scores.most_common(k)

    def keywords(self, doc_tokens, doc_id, k=10):
        weights = {term: self.term_weight(term, doc_id) for term in dict.fromkeys(doc_tokens)}
        # Sort by weight, and by term for ties to make the order reproducible.
        return sorted(weights.items(), key=lambda x: (-x[1], x[0]))[:k]

keywords() ranks the terms of a document by their BM25 weights. A term gets a high weight when it appears frequently in the document but rarely in the other documents, so the top terms characterize the document.

search() adds up the weights of the query terms for each document. Only the documents that contain a query term are scored, so a query without any known term returns an empty list.

Run the script

Run the following command to tokenize the corpus, extract keywords and search documents. Tokenizing all articles takes several minutes.

python tutorial_bm25.py

tutorial_bm25.py

import io
import math
import os
import tarfile
import urllib.request
from collections import Counter

import nagisa
from tqdm import tqdm

CORPUS_URL = "https://www.rondhuit.com/download/ldcc-20140209.tar.gz"
CORPUS_DIR = "ldcc"

# Content words used as index terms
TARGET_POSTAGS = ["名詞", "動詞", "形容詞", "形状詞", "英単語"]


def download_corpus():
    if os.path.isdir(CORPUS_DIR):
        return

    with urllib.request.urlopen(CORPUS_URL, timeout=120) as resp:
        data = resp.read()

    with tarfile.open(fileobj=io.BytesIO(data), mode="r:*") as tar:
        # The filter argument is not available in older Python versions.
        if hasattr(tarfile, "data_filter"):
            tar.extractall(path=".", filter="data")
        else:
            tar.extractall(path=".")
    os.rename("text", CORPUS_DIR)


def load_corpus():
    titles = []
    texts = []
    for category in sorted(os.listdir(CORPUS_DIR)):
        cat_path = os.path.join(CORPUS_DIR, category)
        if not os.path.isdir(cat_path):
            continue
        for fname in sorted(os.listdir(cat_path)):
            if not fname.endswith(".txt") or fname.startswith("LICENSE"):
                continue
            with open(os.path.join(cat_path, fname), encoding="utf-8") as f:
                lines = f.read().splitlines()

            # line 1: URL, line 2: date, line 3: title, line 4-: body
            title = lines[2]
            body = "\n".join(lines[3:])
            titles.append(title)
            texts.append(title+"\n"+body)
    return titles, texts


def split_words(text):
    # nagisa works on a sentence, so tokenize the text line by line.
    words = []
    for line in text.splitlines():
        tokens = nagisa.extract(line, extract_postags=TARGET_POSTAGS)
        words += [word.lower() for word in tokens.words]
    return words


def tokenize(text):
    return [word for word in split_words(text) if word not in nagisa.stopwords]


class BM25:

    def __init__(self, k1=1.5, b=0.75):
        self.k1 = k1
        self.b = b

    def fit(self, corpus_tokens):
        self.num_docs = len(corpus_tokens)
        self.doc_lens = [len(tokens) for tokens in corpus_tokens]
        self.avgdl = sum(self.doc_lens) / self.num_docs

        # Inverted index: term -> {doc_id: term frequency}
        self.index = {}
        for doc_id, tokens in enumerate(corpus_tokens):
            for term, tf in Counter(tokens).items():
                self.index.setdefault(term, {})[doc_id] = tf

        # Inverse document frequency of each term
        self.idf = {}
        for term, postings in self.index.items():
            df = len(postings)
            self.idf[term] = math.log(1 + (self.num_docs - df + 0.5) / (df + 0.5))
        return self

    def term_weight(self, term, doc_id):
        tf = self.index.get(term, {}).get(doc_id, 0)
        if tf == 0:
            return 0.0
        norm = self.k1 * (1 - self.b + self.b * self.doc_lens[doc_id] / self.avgdl)
        return self.idf[term] * tf / (tf + norm)

    def search(self, query_tokens, k=5):
        # Only documents that contain a query term get a score.
        scores = Counter()
        for term in query_tokens:
            for doc_id in self.index.get(term, {}):
                scores[doc_id] += self.term_weight(term, doc_id)
        return scores.most_common(k)

    def keywords(self, doc_tokens, doc_id, k=10):
        weights = {term: self.term_weight(term, doc_id) for term in dict.fromkeys(doc_tokens)}
        # Sort by weight, and by term for ties to make the order reproducible.
        return sorted(weights.items(), key=lambda x: (-x[1], x[0]))[:k]


def main():
    download_corpus()
    titles, texts = load_corpus()
    corpus_tokens = [tokenize(text) for text in tqdm(texts)]

    bm25 = BM25().fit(corpus_tokens)
    print("documents: {}, vocabulary: {}".format(bm25.num_docs, len(bm25.index)))

    # Keyword extraction
    for doc_id in [1100, 4058, 6233]:
        print("\n[{}] {}".format(doc_id, titles[doc_id]))
        keywords = bm25.keywords(corpus_tokens[doc_id], doc_id, k=10)
        print("  "+", ".join(term for term, weight in keywords))

    # Document search
    for query in ["iPhoneの便利なアプリ", "サッカー日本代表の試合"]:
        print("\nquery: {}".format(query))
        for doc_id, score in bm25.search(tokenize(query), k=5):
            print("  {:.2f}  {}".format(score, titles[doc_id]))


if __name__ == "__main__":
    main()

This is the output of tutorial_bm25.py.

documents: 7367, vocabulary: 68012

[1100] スマホでカーナビを実現！純正ドリンクホルダーがスマホホルダーに早変わり【イケショップのレア物】
  ホルダー, holder, ボールジョイント, ドリンク, car, かぶせる, mount, 両側, ミリ, 固定

[4058] 大人にふさわしいビターな味わい　ショコラケーキで聖夜を飾る
  引き渡し, ビスキュイ, ショコラ, ビター, ヘネシー, ケーキ, ガナッシュ, チョコレート, tel, 濃厚

[6233] 香川真司、メッシを語る「分かっていても止められない」
  メッシ, シャビ, リオネル, バルセロナ, イニエスタ, マドリード, 切り返し, トリオ, ミラン, レアル

query: iPhoneの便利なアプリ
  5.55  ようやくiPhoneでもChromeが使えるように！【iPhoneでチャンスを掴め】
  5.31  もっさりiPhoneとはさよなら！ワンタップでサクサクにできる必須アプリ「もっさり解消 for iPhone」【iPhoneアプリ】
  5.09  面倒な画面の明るさ設定はiPhoneを傾けるだけで好みの明るさに！「Tilt de Brightness(明るさ調整)」【iPhoneアプリ】
  5.02  二画面の便利さを450円で！「iDisplay」でiPadをパソコンのセカンドディスプレイ化【iPhoneアプリ】【iPadアプリ】
  5.01  女子だらけの公開イベントに潜入！iPhone女子部がKDDIデザイニングスタジオで「公開リアル部室」を開催【レポート】

query: サッカー日本代表の試合
  7.56  代表戦皮肉るセルジオ越後氏、「これこそザ・フレンドリーマッチ」
  7.41  U-23日本代表、まさかのドロー。セルジオ氏は「まるで高校サッカーの決勝」
  7.35  日本代表戦を速報する中西哲生さん、新幹線から粘りのツイート
  7.17  日本代表新アウエーユニに「こっちをホームにしろよ」の声
  7.17  【五輪サッカー】 韓国に負けたイギリス代表に英ネットユーザーから批判相次ぐ

Use BM25S (BM25-Sparse) with nagisa

BM25S is a fast BM25 library that computes the scores of all documents in advance and stores them in a sparse matrix. Its default tokenizer splits text at non-word characters such as spaces and punctuation.

import bm25s

print(bm25s.tokenize(["日本語の文書を検索します。"], return_ids=False))
#=> [['日本語の文書を検索します']]

You can use nagisa by passing split_words() to the splitter argument of bm25s.tokenization.Tokenizer, and nagisa.stopwords to the stopwords argument. The following script builds the index with the same tokens as tokenize() above, saves the index and the vocabulary of the tokenizer, loads them again and searches documents. It imports functions from tutorial_bm25.py, so put both scripts in the same directory and run it there. It also tokenizes the whole corpus, which takes several minutes.

python tutorial_bm25s.py

tutorial_bm25s.py

import bm25s
import nagisa

from tutorial_bm25 import download_corpus, load_corpus, split_words

INDEX_DIR = "ldcc_bm25s_index"


def build_tokenizer():
    # split_words() lowercases the words after tokenization, so set lower=False.
    return bm25s.tokenization.Tokenizer(lower=False, splitter=split_words,
                                        stopwords=nagisa.stopwords)


def main():
    download_corpus()
    titles, texts = load_corpus()

    # The default tokenizer of bm25s cannot split Japanese text into words.
    print(bm25s.tokenize(["日本語の文書を検索します。"], return_ids=False, show_progress=False))

    # Create the index with the nagisa tokenizer
    tokenizer = build_tokenizer()
    corpus_tokens = tokenizer.tokenize(texts, return_as="tuple")
    retriever = bm25s.BM25(k1=1.5, b=0.75)
    retriever.index(corpus_tokens)

    # Save the index with the titles, and the vocabulary of the tokenizer
    retriever.save(INDEX_DIR, corpus=titles)
    tokenizer.save_vocab(INDEX_DIR)

    # Load them again and search documents
    retriever = bm25s.BM25.load(INDEX_DIR, load_corpus=True)
    tokenizer = build_tokenizer()
    tokenizer.load_vocab(INDEX_DIR)

    queries = ["iPhoneの便利なアプリ", "サッカー日本代表の試合"]
    query_tokens = tokenizer.tokenize(queries, update_vocab=False)
    results, scores = retriever.retrieve(query_tokens, k=5)

    for query, docs, doc_scores in zip(queries, results, scores):
        print("\nquery: {}".format(query))
        for doc, score in zip(docs, doc_scores):
            print("  {:.2f}  {}".format(score, doc["text"]))


if __name__ == "__main__":
    main()

This is the output of tutorial_bm25s.py.

[['日本語の文書を検索します']]

query: iPhoneの便利なアプリ
  5.55  ようやくiPhoneでもChromeが使えるように！【iPhoneでチャンスを掴め】
  5.31  もっさりiPhoneとはさよなら！ワンタップでサクサクにできる必須アプリ「もっさり解消 for iPhone」【iPhoneアプリ】
  5.09  面倒な画面の明るさ設定はiPhoneを傾けるだけで好みの明るさに！「Tilt de Brightness(明るさ調整)」【iPhoneアプリ】
  5.02  二画面の便利さを450円で！「iDisplay」でiPadをパソコンのセカンドディスプレイ化【iPhoneアプリ】【iPadアプリ】
  5.01  女子だらけの公開イベントに潜入！iPhone女子部がKDDIデザイニングスタジオで「公開リアル部室」を開催【レポート】

query: サッカー日本代表の試合
  7.56  代表戦皮肉るセルジオ越後氏、「これこそザ・フレンドリーマッチ」
  7.41  U-23日本代表、まさかのドロー。セルジオ氏は「まるで高校サッカーの決勝」
  7.35  日本代表戦を速報する中西哲生さん、新幹線から粘りのツイート
  7.17  日本代表新アウエーユニに「こっちをホームにしろよ」の声
  7.17  【五輪サッカー】 韓国に負けたイギリス代表に英ネットユーザーから批判相次ぐ

The results are the same as those of tutorial_bm25.py. The scores can differ slightly in the last digits because BM25S uses float32.

Notes on using BM25S with nagisa:

• Set lower=False. By default, BM25S lowercases the text before splitting it, but nagisa may split and tag lowercased text differently from the original text. split_words() lowercases the words after tokenization instead.

• Pass return_as="tuple" when tokenizing the corpus so that the index stores the vocabulary. Otherwise, queries given as lists of strings get a score of 0 for all documents.

• Save the vocabulary of the tokenizer with tokenizer.save_vocab() and load it with tokenizer.load_vocab(). A new tokenizer has an empty vocabulary, so tokenize(queries, update_vocab=False) returns no tokens without it.

• Unlike search() above, BM25S returns k documents even if a query has no known terms. In that case, all the scores are 0, so ignore the results with a score of 0.