JavaScript でやってるのを見かけたので。
cf. JavaScriptで全文検索(N-gram)を実装してみる！ – Simple is Beautiful.
アルゴリズムは N-gram っていう方法のうちでも簡単な uni-gram っていうみたい。詳しくはリンク先の記事を見て。
記事の説明を読んで、なんとなく理解したので Python で書いてみた。リンク先の JavaScript の実装はファイル名を参考にしたくらいで、ちゃんと読んでない。
できたのはこんな感じ。
- create_index.py でインデックスを作って、search.py で検索
- 検索対象のファイルはテキストファイルのみ。documents ディレクトリに入ってる
- インデックスファイルは indexes ディレクトリ内につくる
- document ID とファイル名の対応は docs.json ファイル
インデックスを作る create_index.py
from unigram import document
import json
import os
DOC_DIR = 'documents'
INDEX_DIR = 'indexes'
DOC_DATA = 'docs.json'
files = [os.path.join(DOC_DIR, f) for f in os.listdir(DOC_DIR)]
docs = {}
doc_id = 0
for file in files:
with open(file, 'r') as f:
text = f.read()
tokens = document.tokenize(text)
index = document.classify(tokens)
document.save_index(index, doc_id, INDEX_DIR)
docs[str(doc_id)] = {'name': os.path.basename(file), 'path': file}
doc_id += 1
with open(DOC_DATA, 'w') as f:
json.dump(docs, f, indent=2)
検索コマンド search.py
from unigram import document
import json
import os
import sys
INDEX_DIR = 'indexes'
DOC_DATA = 'docs.json'
string = sys.argv[1]
with open(DOC_DATA, 'r') as f:
docs = json.load(f)
fcount = len(docs)
index_files = list(map(lambda x: os.path.join(INDEX_DIR, x), os.listdir(INDEX_DIR)))
index = {}
for file in index_files:
c = chr(int(os.path.basename(file).replace('.index', '')))
with open(file, 'r') as f:
index[c] = document.parse_index(f.read())
m = list(map(lambda x: index[x], list(string)))
for i in range(fcount):
doc_id = str(i)
if not all(map(lambda x: doc_id in x.keys(), m)):
continue
n = list(map(lambda x: x[doc_id], m))
s = set(n[0])
for s1 in n[1:]:
s = set(list(map(lambda x: x + 1, s)))
s = s & set(s1)
if len(s) > 0:
pos = list(map(lambda x: x - len(string) + 1, s))
pos.sort()
print(docs[doc_id]['name'], pos)
両方から使うモジュール unigram/document.py
import os
def tokenize(text):
return list(text)
def classify(token_list):
tokens = {}
pos = 0
for t in token_list:
if not t in tokens:
tokens[t] = []
tokens[t].append(pos)
pos += 1
return tokens
def save_index(index, doc_id, index_dir):
for c, idx in index.items():
l = str(doc_id) + ':' + ','.join(list(map(lambda x: str(x), idx))) + '\n'
with open(os.path.join(index_dir, str(ord(c)) + '.index'), 'a') as f:
f.write(l)
def parse_index(content):
l = content.split('\n')
l.pop()
index = {}
for l2 in l:
a = l2.split(':')
index[a[0]] = [int(x) for x in a[1].split(',')]
return index
とにかくまずは動くものを、ってことで作ったので、コードが整理されてないのは目を瞑ってほしい（後で直す）。
検索対象ファイルのサンプルには、別のプロジェクトで書いた Ruby のソースファイル。
takatoh@apostrophe $ ls documents Gemfile Rakefile boot.rb config.yaml Gemfile.lock app.rb config.ru config.yaml.example
インデックスを作る。
takatoh@apostrophe $ python create_index.py
できたインデックスファイルがこれ。
takatoh@apostrophe $ ls indexes 10.index 111.index 123.index 44.index 58.index 71.index 84.index 100.index 112.index 124.index 45.index 60.index 72.index 85.index 101.index 113.index 125.index 46.index 61.index 73.index 87.index 102.index 114.index 126.index 47.index 62.index 74.index 89.index 103.index 115.index 32.index 48.index 63.index 75.index 91.index 104.index 116.index 34.index 49.index 64.index 76.index 93.index 105.index 117.index 35.index 50.index 65.index 77.index 95.index 106.index 118.index 36.index 51.index 66.index 78.index 97.index 107.index 119.index 39.index 52.index 67.index 79.index 98.index 108.index 120.index 40.index 53.index 68.index 80.index 99.index 109.index 121.index 41.index 55.index 69.index 82.index 110.index 122.index 43.index 56.index 70.index 83.index
「require」という文字列を検索してみる。ファイル名と出現位置（のリスト）が出力される。
takatoh@apostrophe $ python search.py require boot.rb [0, 17] Rakefile [0, 15, 32, 71] config.ru [0] app.rb [0, 23, 40, 56, 71]
大丈夫そうだ。