最近、簡単な文章検索を実装する機会があったので記事にしました。

githubのほうにプログラム全体載せてあるのでそちらも参考にしてください。
https://github.com/mygod877/cosine_similarity

概要

簡単に言うと、入力した文章(query)で３０件程度の論文のアブストラクトから
最もqueryにマッチしたものを抽出するというものです。

このとき、queryと論文アブストがどれだけ適切か、という指標が必要になります。
今回はこの値にコサイン類似度を用います。

処理フロー

1. queryを受け取る
2. queryと論文アブストの文章を単語ごとに区切り、それぞれTF-IDF法でベクトル化する
3. queryとそれぞれの論文アブストのベクトル間のコサイン類似度を算出する
4. コサイン類似度が最も大きいものを表示

TF-IDF法だとか、コサイン類似度だとか、なにやら難しい単語が続きますが
Pythonのscikit-learnというライブラリで全て勝手にやってくれるので
プログラムとしては簡単なものとなっています。

プログラム

import argparse
import numpy as np
import glob
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity


def main(args):
    data = get_data(args.datafile)
    cs = get_cs(args.query, data)

    max_index = np.argmax(cs)
    max_cs = cs[max_index][0]
    max_data = data[max_index]

    if max_cs > 1e-10:
        print(f"コサイン類似度: {max_cs}")
        print(f"アブストラクト: '{max_data}'")
    else:
        print("NotFound")

def get_data(datafile):
    abstract = np.loadtxt(f"{datafile}", encoding="utf-8", delimiter='|', dtype=str)
    data = []
    for s in abstract:
        text = s.replace(".", "")
        text = text.replace(",", "")
        text = text.replace("(", "")
        text = text.replace(")", "")
        text = text.replace("-", " ")
        data.append(text.lower())

    return data

def get_cs(query, data):
    tfidf = TfidfVectorizer()
    abstract_vector = tfidf.fit_transform(data).toarray()
    query_vector = tfidf.transform([query]).toarray()
    cs = cosine_similarity(abstract_vector, query_vector)
    
    return cs

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("query", type=str)
    parser.add_argument("--datafile", type=str, default="./data/abstract.txt")
    args = parser.parse_args()
    main(args)