文書の類似度判定(3)

データベースに文書を保存するまでが出来たので、やっと本題の類似度判定を。

データベースから文書を取り出して学習と判定を行う

# -*- coding: utf-8 -*-
# ------------------------------------------------------------------ import(s)
import sys
import os
import sqlite3
import gensim.models.doc2vec
import gensim.models.doc2vec


# ------------------------------------------------------------------- const(s)
# ------------------------------------------------------------------- class(s)
# ---------------------------------------------------------------- function(s)
def main():

    list_document = []

    o_conn = sqlite3.connect("dccol.db")
    o_cursor = o_conn.cursor()

    o_cursor.execute(
        """
        SELECT
            id_doc
        ,   word_pos
        ,   word_surface
        ,   word_feature
        FROM
            doc_collection
        WHERE
            LENGTH(word_surface) > 1
            AND
            word_feature = '名詞'
        ORDER BY
            id_doc, word_pos
        """
    )

    dict_document = {}

    for r in o_cursor.fetchall():
        id_doc = r[0]
        word_pos = r[1]
        word_surface = r[2]
        word_feature = r[3]

        if id_doc not in dict_document:
            dict_document[id_doc] = []
        dict_document[id_doc].append(word_surface.encode("utf-8"))

    list_tagged_document = []

    for id_doc, list_word in dict_document.items():

        with open(id_doc + ".wakati", "w") as h_writer:
            h_writer.write(" ".join(list_word))

        o_doc = gensim.models.doc2vec.TaggedDocument(
            words=list_word,
            tags=[id_doc]
        )

        list_tagged_document.append(o_doc)

    dvec_model = gensim.models.doc2vec.Doc2Vec(
        documents=list_tagged_document,
        min_count=1
    )

    list_word = ["ヴィヴィオ", "アインハルト","コロナ","リオ"]
    for similar_word in list_word:
        print u"指定した単語に関連するもの", similar_word.decode("utf-8")
        for r in dvec_model.wv.most_similar(positive=[similar_word]):
            print "\t", r[0].decode("utf-8"), r[1]
        print

    similar_doc = os.path.join("doc_source", "doc1.txt")
    print u"指定した文書に近いもの", similar_doc
    for r in dvec_model.docvecs.most_similar(positive=[similar_doc]):
        print "\t", r[0], r[1]
    print

if __name__ == "__main__":
    main()

# [EOF]

やっていることは、データベースから二文字以上の名詞を取り出して学習用データを生成して、与えたキーワード、文書に似ている（近いベクトルのもの）を取得しています。

実行すると以下のような結果が表示されます。

指定した単語に関連するもの ヴィヴィオ
        スタイル 0.404163241386
        状態 0.396886348724
        ため 0.391772806644
        アインハルト 0.378925800323
        リンネ・ベルリネッタ 0.374079525471
        結果 0.3656296134
        意味合い 0.356572329998
        デバイス 0.350234866142
        雪辱 0.344531387091
        兵器 0.343862652779

学習させた文書は、Pixiv百科事典からキャラクター名の項目をテキスト化したものを与えています。

ヴィヴィオに関連するものとして、

• アインハルト
• リンネ・ベルリネッタ

といったキーワードが出力されています。

与えた文書量が少なく、文書自体も短いためなんとなくうまくいっているようないないような。