オレオレ RAG をさくっと作る

この記事は間違いが含まれている可能性があります。

もともと自社のドキュメントでは Meilisearch で日本語全文検索を実現していましたが、ドキュメントに質問できるようしたいと思い、簡単な RAG を作りたい！と思っていました。

とりあえず、ドキュメントを分割し、ベクトル化してベクトルデータベースに突っ込んで、質問をベクトル化して結果を引っ張り、それを LLM に食べさせて解説させる、というのができればよいということがわかりました。

ベクトル化はよく使われている OpenAI Embeddings API を利用し、ベクトルデータベースは普段からよく使っている DuckDB に VSS (Vector Similarity Search for DuckDB) という拡張を使うことにしました。

自社のドキュメントをなんとかうまいこと分割して、あとは分割したドキュメントを API を叩いてベクトル化して、 DuckDB に保存しました。OpenAI Embeddings API は想像以上にとても安いので自社製品のドキュメント程度のサイズであれば 5 円もかかりませんでした。

実際に色々動かしてみて、ハルシネーションをほぼ感じない出力になりました。ただ OpenAI Embeddings API は従量課金なので、クエリーを実行するたびに毎回ベクトル化するのは、誰もが利用できるサービスとして利用するには、少し抵抗がありました。

そんなときに 株式会社Preferred Networks さんから Plamo-Embedding-1B という「日本語に強い」ベクトル化のモデルが公開されました。OSS でライセンスは Apache-2.0 です。さらに、ローカルの CPU でさくさくっと動きびっくりしました。

以下は実際にどんなものかをイメージしやすくするための pyproject.toml と main.py です。

[project]
name = "oreore-rag"
version = "0.1.0"
readme = "README.md"
requires-python = ">=3.11"
dependencies = [
    "duckdb>=1.2.2",
    "sentencepiece>=0.2.0",
    "torch>=2.6.0",
    "transformers>=4.51.3",
]

import duckdb
import torch
from transformers import AutoModel, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained(
    "pfnet/plamo-embedding-1b", trust_remote_code=True
)
model = AutoModel.from_pretrained("pfnet/plamo-embedding-1b", trust_remote_code=True)

device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device)

# https://sora.shiguredo.jp/
docs = [
    "WebRTC による音声・映像・メッセージメッセージのリアルタイムな配信と、その録音・録画を実現します",
    "お客様ご自身のサーバーにインストールしてご利用いただくパッケージソフトウェアです",
    "株式会社時雨堂がフルスクラッチで開発しており、日本語によるサポートとドキュメントを提供します",
]


def main():
    # メモリ上にデータベースを作成
    conn = duckdb.connect()
    conn.sql("INSTALL vss")
    conn.sql("LOAD vss")

    conn.sql("CREATE SEQUENCE IF NOT EXISTS id_sequence START 1;")
    conn.sql(
        "CREATE TABLE IF NOT EXISTS sora_doc (id INTEGER DEFAULT nextval('id_sequence'), content TEXT, vector FLOAT[2048]);"
    )

    with torch.inference_mode():
        for doc, doc_embedding in zip(docs, model.encode_document(docs, tokenizer)):
            conn.execute(
                "INSERT INTO sora_doc (content, vector) VALUES (?, ?)",
                [
                    doc,
                    doc_embedding.cpu().squeeze().numpy().tolist(),
                ],
            )

    query = "時雨堂について教えてください"
    print("query:", query)

    with torch.inference_mode():
        query_embedding = model.encode_query(query, tokenizer)
        result = conn.sql(
            """
            SELECT content, array_cosine_distance(vector, ?::FLOAT[2048]) as distance
            FROM sora_doc
            ORDER BY distance
            """,
            params=[query_embedding.cpu().squeeze().numpy().tolist()],
        )

        for row in result.fetchall():
            print("distance:", row[1], "|", row[0])


if __name__ == "__main__":
    main()
    # query: 時雨堂について教えてください
    # distance: 0.22656434774398804 | 株式会社時雨堂がフルスクラッチで開発しており、日本語によるサポートとドキュメントを提供します
    # distance: 0.39890891313552856 | お客様ご自身のサーバーにインストールしてご利用いただくパッケージソフトウェアです
    # distance: 0.5286199450492859 | WebRTC による音声・映像・メッセージメッセージのリアルタイムな配信と、その録音・録画を実現します

たったこれだけで RAG が実現できます。もちろん色々チューニングの余地はあると思いますが自分が使う分だけなら十分だと思います。

あとは MCP で繋げられるようにすれば、VS Code から気軽に利用できますので、ぜひ作ってみてください。