埋め込みモデルによる再評価

RAG のrerankingに関するPythonコード

目次

Reranking は、Retrieval Augmented Generation (RAG) システム における 2 番目のステップであり、Retrieval（検索）と Generation（生成）の間に位置します。

Electric cubes in digital space

上記は、Flux-1 dev が「デジタル空間内の電気立方体」をどのようにイメージしているかを示しています。

RAG システム構築の完全なガイドについては、Retrieval-Augmented Generation (RAG) Tutorial: Architecture, Implementation, and Production Guide をご覧ください。

この記事では、埋め込みモデルを用いた再ランク付けについて解説します。代替アプローチとして、Reranking texts with Ollama and Qwen3 Embedding LLM - in Go や、Reranking documents with Ollama and Qwen3 Reranker model - in Go もご参照いただけます。

再ランク付けによる検索

最初にベクター DB に埋め込み形式でドキュメントを保存しておけば、検索処理ですぐに類似ドキュメントのリストが取得できます。

スタンアロンの再ランク付け

しかし、まずインターネットからドキュメントをダウンロードする場合、検索システムの応答は、検索プロバイダーの好みやアルゴリズム、スポンサー付きコンテンツ、SEO 最適化などに影響を受ける可能性があります。そのため、検索後の再ランク付け（post-search reranking）が必要です。

私が採用した方法は以下の通りです。

  • 検索クエリの埋め込みを取得する
  • 各ドキュメントの埋め込みを取得する（ドキュメントは 8k トークンを越えることは想定されていません）
  • クエリと各ドキュメントの埋め込み間の類似度を計算する
  • この類似度に基づいてドキュメントをソートする

ここではベクター DB を使用しません。さっそく始めましょう。

サンプルコード

Langchain を使用して Ollama に接続し、langchain の cosine_similarity 関数を利用しています。 類似度の閾値は、ドメインや埋め込み LLM によって異なるため、注意が必要です。

このコードが何らかの形で役立つことを願っています。 Copy/Paste/UseAnyWayYouWant ライセンスです。 どうぞよろしくお願いいたします。

from langchain_core.documents import Document
from langchain_community.embeddings import OllamaEmbeddings
from langchain_community.utils.math import cosine_similarity
import numpy as np


def cosine_distance(a: np.ndarray, b: np.ndarray) -> np.ndarray:
    return 1.0 - cosine_similarity(a, b)

def compute_score(vectors: np.ndarray) -> float:
    score = cosine_distance(vectors[0].reshape(1, -1), vectors[1].reshape(1, -1)).item()
    return score

def list_to_array(lst):
    return np.array(lst, dtype=float)   

def compute_scorel(lists) -> float:
    v1 = list_to_array(lists[0])
    v2 = list_to_array(lists[1])
    return compute_score([v1, v2])

def filter_docs(emb_model_name, docs, query, num_docs):
    content_arr = [doc.page_content for doc in docs]

    ollama_emb = OllamaEmbeddings(
        model=emb_model_name
    )

    docs_embs = ollama_emb.embed_documents(content_arr)
    query_embs = ollama_emb.embed_query(query)
    sims = []
    for i, emb in enumerate(docs_embs):
        idx = docs[i].id
        s = compute_scorel([query_embs, docs_embs[i]])
        simstr = str(round(s, 4))
        docs[i].metadata["sim"] = simstr
        sim = {
            "idx": idx,
            "i": i,
            "sim": s,
        }
        sims.append(sim)

    sims.sort(key=sortFn)

    sorted_docs = [docs[x["i"]] for x in sims]
    filtered_docs = sorted_docs[:num_docs]
    return filtered_docs

最適な埋め込みモデル

私のタスクにおいて、現時点で最も優れた埋め込みモデルは bge-large:335m-en-v1.5-fp16 です。

2 位には nomic-embed-text:137m-v1.5-fp16 と jina/jina-embeddings-v2-base-en:latest が挙げられます。

ただし、ご自身のドメインやクエリに合わせて独自にテストを行うことをお勧めします。

関連リンク