転置インデックスによる検索システムを作ってみた
転置インデックスによる検索システムを作ってみよう！ にインスパイアされて作ってみました。
検索記事は
[記事ID][SPC][記事内容]\n
検索対象ファイルとして
1 これはペンです 2 最近はどうですか？ 3 ペンギン大好き 4 こんにちは。いかがおすごしですか？ 5 ここ最近疲れ気味 6 ペンキ塗りたてで気味が悪いです 7 ペンペンペンペン
という内容のtest.txt用意しました。
インデックス
n-gramをkeyとして、各記事のtf(記事中のn-gram出現頻度)と記事IDのタプルをtf降順にsortしたリストを登録した辞書
index[n-gram] => [(tf, 記事ID), ...] #タプルはtf降順にsortしておく
をcPickleでシリアライズしたものをインデックスファイルとして使うことにします。
cPickleを使うと
Pythonのオブジェクト -> pickle -> 文字列
という変換をして、変換した文字列から
文字列 -> pickle -> Pythonのオブジェクト
といった具合に元のオブジェクトに戻すことができます。
具体的に書くと
#pickletest.py
import cPickle
a = range(5)
print a
obj = cPickle.dumps(a)
print obj
b = cPickle.loads(obj)
print b
a = {}
for i in range(5):
a[i] = str(i**10)
print a
obj = cPickle.dumps(a)
print obj
b = cPickle.loads(obj)
print b> python pickletest.py
[0, 1, 2, 3, 4]
(lp1
I0
aI1
aI2
aI3
aI4
a.
[0, 1, 2, 3, 4]
{0: '0', 1: '1', 2: '1024', 3: '59049', 4: '1048576'}
(dp1
I0
S'0'
sI1
S'1'
sI2
S'1024'
p2
sI3
S'59049'
p3
sI4
S'1048576'
p4
s.
{0: '0', 1: '1', 2: '1024', 3: '59049', 4: '1048576'}といったようなことが可能です。
以上のことを踏まえると、インデックスを作るプログラム、makeindex.pyを以下のようになります。
#!/usr/bin/python
# -*- coding: utf-8 -*-
import sys
import codecs
import cPickle
# n-gramのn
n = int(sys.argv[1])
# インデックス作成対象ファイルの名前
input_file = sys.argv[2]
index = {}
doc_num = 0
# 文字コードを指定してファイルを開き、一行ずつ読み込む
for line in codecs.open(input_file, 'r', 'utf-8'):
d = {}
# 読み込んだ行を記事IDと記事に分解
l = line.split(' ')
id = int(l[0])
doc = l[1][:-1]
# 記事の文字数回ループ
for i in range(len(doc)):
# 記事のi文字目からn文字をn-gramとする
ngram = doc[i:i+n]
try:
# 既に出現していたら出現回数+1
d[ngram] += 1
except:
# 初めてなら出現回数1
d[ngram] = 1
# 記事に出現したn-gram全てに対して...
for ng in d:
# (tf, 記事ID) というタプルを作る
# tf = n-gramが記事に出現した回数 / 記事の文字数
t = (1.0*d[ng]/len(doc), id)
try:
# 他の記事で既に出現していたら単純に追加
index[ng].append(t)
# tf降順でsort
index[ng].sort(reverse=True)
except:
# 初めて出現したn-gramならリストを作って自分を追加
index[ng] = [t]
# 後でtf-idfを計算するために記事の数を覚えておく
doc_num += 1
# indexに記事の数を登録
index['__DOC_NUM__'] = doc_num
# indexをcPickleで文字列に変換してファイルに書き込む
file(input_file + '-' + str(n) + 'g.idx', 'w').write(cPickle.dumps(index))例えばtest.txtの2gram単位のインデックスファイルを作りたい場合は
> ./makeindex.py 2 test.txt
というように実行することで、test.txt-2g.idxというインデックスファイルが作成されます。
検索
検索を行うプログラム、search.pyは以下のようになります。
2007/12/05追記: 記事のスコアにtf-idfを加算するところを間違えていたので修正しました。すいません……。この修正により、最後の検索結果も変わります。
#!/usr/bin/python
# -*- coding: utf-8 -*-
import sys
import codecs
import cPickle
import math
# インデックスファイルの名前
index_file = sys.argv[1]
# インデックスファイルを読み込み、cPickleで文字列から辞書に変換
index = cPickle.loads(file(index_file).read())
# インデックスファイルの名前からn-gramのnを決定
n = int(index_file.split('-')[-1][0])
# 標準入力の文字コードを指定
sys.stdin = codecs.getreader('utf-8')(sys.stdin)
# 検索結果のスコアを登録する辞書を用意
result = {}
# 標準入力から検索語を読み込む
for line in sys.stdin:
print line,
# 検索語の文字数回ループ
for i in range(len(line)):
# 検索語のi文字目からn文字をn-gramとする
ngram = line[i:i+n]
# n-gramが出現した記事があるなら...
if ngram in index:
# indexから記事数取得
doc_num = index['__DOC_NUM__']
# n-gramが出現した記事のリスト取得
elems = index[ngram]
# df算出
# df = n-gramが出現した記事数 = n-gramが出現した記事のリストの長さ
df = len(elems) + 1
# n-gramが出現した記事のリストに対して...
for tf, id in elems:
# n-gramに対する記事のtf-idf算出
# tf-idf = tf * log(総記事数 / n-gramが出現した記事数)
tfidf = tf * math.log(1.0*doc_num/df)
try:
# 既に他のn-gramでスコアが算出されていたらtf-idfを加算
# 修正前:
# score = resul[id] <- resultのtypo(ここで例外発生)。タプルをそのままscoreとするのも間違い。
# result[id] += (score + tfidf, id) <- タプルにタプルを足している……！
# 2007/12/05修正:
# タプルの最初の要素をスコアとして、tfidfを加算するのが正しい
score = result[id][0]
result[id] = (score + tfidf, id)
except:
# 記事IDのスコアをtf-idfとして登録
result[id] = (tfidf, id)
# 辞書resultを値だけのリストにして、スコア降順にソート
for score, id in sorted(result.values(), reverse=True):
# 結果出力
print 'ID:' + str(id), 'SCORE:' + str(score)基本的には
という流れになっています。
実際に使ってみると以下のようになります。
> echo '最近ペンギンが好き' | ./search.py test.txt-2g.idx 最近ペンギンが好き ID:3 SCORE:0.178966138356 ID:7 SCORE:0.168236118311 ID:5 SCORE:0.105912232548 ID:2 SCORE:0.0941442067097 ID:1 SCORE:0.0480674623745 ID:6 SCORE:0.0224314824414
> echo 'ペンギン' | ./search.py test.txt-2g.idx ペンギン ID:3 SCORE:0.178966138356 ID:7 SCORE:0.168236118311 ID:1 SCORE:0.0480674623745 ID:6 SCORE:0.0224314824414
> echo 'ペン' | ./search.py test.txt-2g.idx ペン ID:7 SCORE:0.168236118311 ID:3 SCORE:0.0480674623745 ID:1 SCORE:0.0480674623745 ID:6 SCORE:0.0224314824414
> echo 'ペ' | ./search.py test.txt-2g.idx ペ
2gram単位のインデックスなので、一文字では検索できませんでした。こういう場合は…… えーっと、どうするのかな……。
2007/12/05追記: '最近ペンギンが好き' で検索しても 'ペンギン' で検索しても 記事ID:3のスコアが同じことから、上の結果はおかしいことがわかります。修正したsearch.pyで検索すると以下のような結果になります。
> echo '最近ペンギンが好き' | ./search.py test.txt-2g.idx 最近ペンギンが好き ID:3 SCORE:0.584965877444 ID:7 SCORE:0.168236118311 ID:5 SCORE:0.105912232548 ID:2 SCORE:0.0941442067097 ID:1 SCORE:0.0480674623745 ID:6 SCORE:0.0224314824414
> echo 'ペンギン' | ./search.py test.txt-2g.idx ペンギン ID:3 SCORE:0.405999739087 ID:7 SCORE:0.168236118311 ID:1 SCORE:0.0480674623745 ID:6 SCORE:0.0224314824414
> echo 'ペン' | ./search.py test.txt-2g.idx ペン ID:7 SCORE:0.168236118311 ID:3 SCORE:0.0480674623745 ID:1 SCORE:0.0480674623745 ID:6 SCORE:0.0224314824414
> echo 'ペ' | ./search.py test.txt-2g.idx ペ
