【TypeScript】N-gramモデルを使ったお手軽全文検索を手抜き実装してみる

はじめに

大量のテキストに対し高速かつ手軽に全文検索を行いたいシチュエーションがあったため、N-gramというモデルを利用して実装してみました。

N-gramとは

自然言語処理におけるN-gramとは、「文字列をN個の連続するシンボルで分割していき、その分割した文字列に対し出現頻度やパターンの分析を行う」という手法です。ここで言うシンボルとは文字や形態素など、自由に決めることができます。

例えば次のような文字列を、Nを2、シンボルを文字とする形で分割してみると、

input:  今日はよい天気です。
output: [今日, 日は, はよ, よい, い天, 天気, 気で, です, す。]

というようになります。

N-gramで全文検索？

検索対象の文字列と検索するクエリ文字列をそれぞれN-gramで分割し、それらを照らし合わせることで関連性の高低を把握できるそうです。

TypeScriptで実装

試しに次のようなPost[]型のデータに対して検索することを考えてみます。

type Post = {
  id: string;
  createdAt: Date;
  text: string; // ここを検索する
  userId: string;
};

まずは検索対象を正規化する関数を作っておきます。大文字・小文字の区別などはなくしておいたほうが検索にはよいでしょう。

const cleanup = (text: string): string => {
  return text.normalize("NFC").toLowerCase();
};

次に、文字列をN-gramで分割する関数を定義します。スプレッド構文で文字ごとに分割したうえで、Array.prototype.reduceを使って分割された文字列の配列とします。（mapとfilterを使っても良かったかもしれません。）

const nGram = (text: string, n: number): string[] => {
  return [...text].reduce<string[]>((acc, _cur, idx, src) => {
    if (idx < n) return acc;
    acc.push(src.slice(idx - n, idx).join(""));
    return acc;
  }, []);
};

それでは、Post[]型の値それぞれを検索しやすいデータ（searchIndex）に変換する関数を定義します。ここでは、分割後文字列それぞれが記事IDとその記事での出現数をもつような型として、Map<分割後文字列, Map<記事ID, 出現数>>を使います。Mapへのアクセスは比較的高速であるため、検索のようなたくさんのアクセスを行う場面では適しているはずです。

const getSearchIndex = (posts: Post[], n: number) => {
  return posts
    // `{ id: 記事ID, value: 分割後文字列 }[]`にする
    .flatMap(({ id, text }) => {
      return nGram(cleanup(text), n).map((value) => ({ id, value }));
    })
    .reduce<Map<string, Map<string, number>>>((acc, { id, value }) => {
      /** 当該分割後文字列で登録済みの情報 */
      const prevValue = acc.get(value);

      if (prevValue === undefined) {
        // 未登録だった場合：記事IDと出現数（1）を登録
        return acc.set(value, new Map([[id, 1]]));
      } else {
        // 登録済みだった場合：記事IDと出現数（+1）を登録
        const prevCount = prevValue.get(id);
        return acc.set(value, prevValue.set(id, (prevCount ?? 0) + 1));
      }
    }, new Map());
};

あとは、検索する文字列に対してもcleanup()とnGram()を適用したうえで、searchIndexに対して問い合わせていくだけです。

const getRanking = (query: string, n: number) => {
  return [
    ...nGram(cleanup(query), n).reduce<Map<string, number>>((acc, cur) => {
      // `searchIndex`から、検索語句（分割後）と対応した`Map<記事ID, 出現数>`を手に入れる
      const result = searchIndex.get(cur);
      if (result === undefined) return acc;

      for (const [postId, count] of result) {
        // ランキングに登録（登録済みだった場合は`count`分を加算）
        const prevCount = acc.get(postId);
        acc.set(postId, (prevCount ?? 0) + count);
      }

      return acc;
    }, new Map()),
  ].sort(([, a], [, b]) => b - a); // 降順ソート
};

これで、検索語句と関連性の高いPostのIDが、関連性が高い順で手に入ります。

おわりに

もう少し頑張って実装するなら、Post["text"]の文字数によってcountの数値を割るような、丁寧なスコア計算をしてもよいかもしれません。この実装では、文字数が多い記事が有利になってしまいます。とはいえ予想よりしっかりと動いてくれたのでひとまずよしとします。