TF-IDF埋め込みとコサイン検索

目標

TF-IDF埋め込みをnumpyで直接計算し、コサイン類似度で文書を検索して、次元削減の代償を数値で確認します。

なぜ重要なのか

埋め込みAPIを呼び出すと、ベクトルが出てきます。その中で何が起きているかを知らないままでも使えますが、検索結果がおかしいときに、原因を絞り込めません。

TF-IDFは、ニューラルネットワークの埋め込みよりはるかに単純ですが、核心の構造は同じです。テキストをベクトルに変え、長さを正規化し、内積で類似度を測るという骨組みが同じです。そのため、ここで正規化を忘れると、長い文書がすべてのクエリに引っかかる現象を直接見られ、その経験が、実際のベクトル検索をデバッグするときに、そのまま役に立ちます。

最後のステップのハッシングトリックにも意味があります。次元を減らすとメモリと速度が得られますが、衝突のために情報を失います。その損失が実際にどれくらいかを数値で測ってみると、ベクトルの次元を選ぶときの感覚が身に付きます。

ステップ

作業ディレクトリは/root/llmです。対象は、docsテーブルの文書30件です。

各文書をトークン化して/root/llm/doc_tokens.tsvに保存してください。docs.id<탭>공백으로 구분한 토큰들の形式で(プレースホルダーはタブと、空白で区切ったトークンです)、30行です。トークン化のルールは、小文字に変換してから、正規表現[가-힣a-z0-9]+で抜き出すことです(この正規表現にはハングルの範囲が含まれます)。
語ごとの文書頻度を/root/llm/df.tsvに항<탭>dfの形式で保存してください(プレースホルダーは語とタブです)。語は昇順にソートし、1つの文書に何回出てきても1と数えます。
IDFを/root/llm/idf.tsvに항<탭>idfの形式で保存してください(プレースホルダーは語とタブです)。語の順序はステップ2と同じで、値は小数点以下6桁です。式はln((1 + N) / (1 + df)) + 1で、Nは文書数です。
TF-IDF行列を/root/llm/tfidf.npyとして保存してください。形は(文書数, 語数)で、語の順序はステップ2と同じです。値は、出現回数にidfを掛けてから、各行のL2の長さを1に正規化したものです。
各文書の最近傍を/root/llm/sim_top.tsvにdocs.id<탭>가장 비슷한 docs.id<탭>점수の形式で保存してください(プレースホルダーはタブと、最も似ているdocs.id、スコアです)。自分自身は除き、スコアは小数点以下6桁で、同点ならdocs.idが小さいほうを選びます。
下のクエリ3つについて、上位3つの文書を/root/llm/query_top3.tsvに질의번호<탭>순위<탭>docs.id<탭>점수の形式で保存してください(プレースホルダーはクエリ番号、順位、スコアと、その間のタブです)。合計9行です。
- クエリ1: 인덱스가 왜 안 타는지 알고 싶다(韓国語で「インデックスがなぜ使われないのか知りたい」を意味するクエリ文です)
- クエリ2: TCP 연결이 안 될 때 무엇을 보나(韓国語で「TCP接続がつながらないとき、何を見るか」を意味するクエリ文です)
- クエリ3: GPU 는 왜 메모리 때문에 느려지나(韓国語で「GPUはなぜメモリのせいで遅くなるのか」を意味するクエリ文です) クエリベクトルも、同じトークン化と同じidfを使い、L2正規化します。同点なら、docs.idが小さいほうが先になります。
ハッシングトリックで1024次元の行列を作って/root/llm/hashed.npyとして保存してください。バケットはint(md5(항).hexdigest(), 16) % 1024(プレースホルダーは語です)で、各バケットに등장 횟수 곱하기 idf(韓国語で「出現回数かけるidf」を意味する表記です)を累積してから、行をL2正規化します。
2つの方式の最近傍が一致する文書数と割合を、/root/llm/compare.tsvに1行で일치수<탭>비율の形式で保存してください(プレースホルダーは一致数、割合と、その間のタブです)。割合は小数点以下3桁です。

参考

numpyはインストールされています。np.save、np.load、np.linalg.normを使います。
文書と語の順序を一度決めたら、最後まで同じ順序を保つ必要があります。
よくある間違い1: Python組み込みのhash()は、実行のたびに値が変わります。必ずhashlib.md5を使ってください。
よくある間違い2: 正規化を忘れると、長い文書がすべてのクエリに引っかかります。各行の長さが1であることを確認してください。

文書をトークン化する

各文書をトークン化して/root/llm/doc_tokens.tsvに保存してください。docs.id<탭>공백으로 구분한 토큰들の形式で(プレースホルダーはタブと、空白で区切ったトークンです)、30行です。トークン化のルールは、小文字に変換してから、正規表現[가-힣a-z0-9]+で抜き出すことです(この正規表現にはハングルの範囲が含まれます)。

小文字に変換してから、ハングル・英字・数字が連続したかたまりだけを残します。正規表現1つで済みます。

文書頻度を数える

語ごとの文書頻度を/root/llm/df.tsvに항<탭>dfの形式で保存してください(プレースホルダーは語とタブです)。語は昇順にソートし、1つの文書に何回出てきても1と数えます。

1つの文書に何回出てきても1と数えます。語は昇順にソートします。

IDFを計算する

IDFを/root/llm/idf.tsvに항<탭>idfの形式で保存してください(プレースホルダーは語とタブです)。語の順序はステップ2と同じで、値は小数点以下6桁です。式はln((1 + N) / (1 + df)) + 1で、Nは文書数です。

分母と分子に1を足し、最後に1を足す、平滑化の式を使います。

TF-IDF行列を作る

TF-IDF行列を/root/llm/tfidf.npyとして保存してください。形は(文書数, 語数)で、語の順序はステップ2と同じです。値は、出現回数にidfを掛けてから、各行のL2の長さを1に正規化したものです。

出現回数にidfを掛けてから、行単位で長さを1にそろえます。順序を変えると値が変わります。