5-16. TF-IDFでポケモンの分類を再定義してみた（ヤドランは”シェルダーポケモン”）

pip install janome

import csv
import math
import numpy as np
import pandas as pd
from janome.tokenizer import Tokenizer

#csvファイルの読み込み、ポケモン名と説明文の辞書型に -> {'フシギダネ': '生まれたときから背中に植物の…'}
data = pd.read_csv('save_151.csv', encoding='UTF-8', low_memory=False)
data = np.array(data)
pokemon_words = dict(zip(data[:, 1], data[:, 2]))
print(pokemon_words['フシギダネ'])

 生まれたときから 背中に 植物の タネが あって 少しずつ 大きく 育つ。 生まれたときから 背中に 不思議な タネが 植えてあって 体と ともに 育つという。 何日だって なにも 食べなくても 元気! 背中のタネに たくさん 栄養が あるから 平気だ! 生まれて しばらくの あいだ 背中の タネに つまった 栄養を とって 育つ。 日なたで 昼寝を する 姿を 見かける。太陽の 光を いっぱい 浴びることで 背中の タネが 大きく 育つのだ。 生まれてから しばらくの あいだは 背中の タネから 栄養を もらって 大きく 育つ。

#janomeで単語に分割する関数
def split_sentence(sentence):
    t = Tokenizer()
    words = []
    for token in t.tokenize(sentence):
        #文書を単語に分割し、動詞、名詞、形容詞のみ抽出する
        if token.part_of_speech.split(',')[0] in ['動詞', '名詞', '形容詞']:
            #()や記号やゲームのタイトルを除外する
            if token.surface in ['(', ')', 'X', 'Y', 'ピカ', 'ブイ', 'ソード', 'シールド', '!', '一']:
                continue
            words.append(token.surface)
    return words

#ポケモンの説明文を分割して単語リストに置き換える -> {'フシギダネ': ['生まれ', 'とき', '背中', ...]}
for name in pokemon_words.keys():
    pokemon_words[name] = split_sentence(pokemon_words[name])  
print(pokemon_words['フシギダネ'])

['生まれ', 'とき', '背中', '植物', 'タネ', 'あっ', '大きく', '育つ', '生まれ', 'とき', '背中', '不思議', 'タネ', '植え', 'あっ', '体', '育つ', 'いう', '何', '日', 'なに', '食べ', '元気', '背中', 'タネ', 'たくさん', '栄養', 'ある', '平気', '生まれ', 'あいだ', '背中', 'タネ', 'つまっ', '栄養', 'とっ', '育つ', '日なた', '昼寝', 'する', '姿', '見かける', '太陽', '光', '浴びる', 'こと', '背中', 'タネ', '大きく', '育つ', 'の', '生まれ', 'あいだ', '背中', 'タネ', '栄養', 'もらっ', '大きく', '育つ']

#TF、単語とtfの辞書を返す
def tf(words):
    result = {}
    for word in words: #同じ単語が２回以上計算される可能性あり（上書きのため結果に影響なし）
        #TF計算
        result[word] = words.count(word) / len(words)
    return result

#IDF、単語とidfの辞書を返す
def idf(words, words_list_noduplicate):
    result = {}
    for word in words: #同じ単語が２回以上計算される可能性あり（上書きのため結果に影響なし）
        #使用されている文書数をカウント
        count = 0
        for name in pokemon_words.keys():
            if word in pokemon_words[name]:
                count += 1
        #IDF計算
        result[word] =  math.log(len(pokemon_words) / count) + 1
    return result

#TF-IDF
def tf_idf(tf_result, idf_result):
    result = {}
    for word in tf_result.keys():
        #TF-IDF計算
        result[word] = tf_result[word] * idf_result[word]
    #降順にソートして上位３件に絞って返す
    return sorted(result.items(), key=lambda x:x[1], reverse=True)[0:3]

results = []

#各ポケモンの処理
for name in pokemon_words.keys():
    #各単語のTF、単語とtfの辞書を返す
    tf_result = tf(pokemon_words[name])
    
    #各単語のIDF、単語とidfの辞書を返す
    idf_result = idf(pokemon_words[name], pokemon_words)
    
    #各単語のTF-IDF
    tf_idf_result = tf_idf(tf_result, idf_result)
    
    #tf_idfの結果リストの先頭にポケモン名を挿入
    tf_idf_result.insert(0, name)
    
    results.append(tf_idf_result)
print(results[0])

#結果の書き出し
with open('tf-idf_151.csv', 'w', encoding='UTF-8', newline='') as f:
    for result in results:
        writer = csv.writer(f)
        writer.writerow(result)

['フシギダネ', ('タネ', 0.4709476754944102), ('背中', 0.33653182736061504), ('育つ', 0.32373349656599193)]