ミニ・ベクトルストアを作る ― 全件検索とtop-k

導入

ここまでの「ベクトル化」と「類似度計算」を1つにまとめ、文書を貯めておいて質問を投げると上位k件を返してくれる、小さな「ベクトルストア」をクラスとして作ってみましょう。これがRAGの**検索器（retriever）**の正体そのものです。

説明

import numpy as np
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity

class MiniVectorStore:
    def __init__(self, docs):
        self.docs = docs
        self.vectorizer = TfidfVectorizer().fit(docs)
        self.matrix = self.vectorizer.transform(docs)

    def search(self, query, k=3):
        q = self.vectorizer.transform([query])
        scores = cosine_similarity(q, self.matrix)[0]
        idx = np.argsort(scores)[::-1][:k]
        return [(self.docs[i], float(scores[i])) for i in idx]

store = MiniVectorStore([
    "Paris is the capital of France.",
    "Tokyo is the capital of Japan.",
    "The Louvre museum is in Paris.",
    "Mount Fuji is near Tokyo.",
    "Python is a programming language.",
])

for doc, score in store.search("famous places in Paris", k=2):
    print(f"{score:.3f}  {doc}")

__init__ でやっているのは、文書を受け取って一度だけベクトル化しておくこと（索引の構築）です。search メソッドが呼ばれるたびに、この索引を毎回作り直す必要はありません。

search の中身は前レッスンとほぼ同じですが、np.argsort(scores) でスコアを小さい順に並べたインデックスを取得し、[::-1] で逆順にして大きい順にし、[:k] で上位k件だけを残しています。これが「top-k検索」と呼ばれる、検索システムの基本パターンです。

実務のRAGシステムでは、ここの TfidfVectorizer の部分を、次のレッスンで紹介する本物の埋め込みモデルに、self.matrix の部分をベクトルDBに置き換えます。仕組みそのものは、このミニ・ベクトルストアと変わりません。

やってみよう

store.search(...) の k を 3 に変えたり、"famous places in Paris" を "Japan mountains" など別の質問に変えたりして、返ってくる結果がどう変わるか確かめてみましょう。

演習

store.search("python programming", k=2) の結果から、文書の文字列だけ（スコアは表示せず）を print で表示してください。