再ランキング

40 件の CLERC 法律クエリに対して 30 パッセージの BM25 候補リストを作り、クエリと候補の各ペアに 1 問の TypeSafe 質問を使って、top-1 精度を 5% から 18% に、top-10 精度を 38% から 62% に引き上げます。

何千もの文書があり、ある特定の質問に答える 1 つを見つける必要があるとします。では、どうやってそれを見つけるのでしょう。

まず、キーワードマッチングのような素早い方法で、何千もの候補をもっともらしい候補の短いリストに絞り込みます。これを高速検索と呼びます。

高速検索はそれに向いていますが、候補リスト上のどの候補が正しいかまでは分かりません。そこで再ランキングの出番です。候補リスト上の各候補をクエリに対して直接スコアリングし、最も良いものを先頭に置きます。

以下の両方のステップを、CLERC データセットの 3,565 件の裁判所意見のパッセージに対して実行します。BM25 が 40 件のクエリそれぞれについて 30 候補の高速検索の候補リストを作り、次に TypeSafe が各候補リストを再ランキングします。再ランキングにより、正しいパッセージが 1 位になるのはクエリの 18% で、高速検索だけの場合の 5% から上昇します。

この過程で、次のことを学びます：

高速検索が何をするか、なぜそれだけでは答えにならないか
再ランキングとは何か、高速検索のステップの後にどう位置づけられるか
TypeSafe がクエリに対して 1 つの候補をどうスコアリングし、それが結果をどれだけ改善するか

自分で試す

TypeSafe プレイグラウンドでクエリ・候補・再ランキングの質問を開く

何千もの文書から 1 つをどう見つけるか

文書の山と、クエリ、つまり探しているものを説明するテキストの断片があるとします。その山のどこかに、それに答える 1 つの文書があります。

すべての文書をクエリと 1 つずつ突き合わせる方法は機能しますが、文書 1 つにつき 1 回の比較で、何百万もの文書はクエリごとに何百万回もの比較を意味します。2 段階のアプローチで性能を改善できます。

山全体に対して実行できるほど速い方法で、山を有望な候補の短いリストに絞り込みます。
その短いリストにより精度の高いステップを適用し、正確な正解を見つけます。

アニメーション図：文書の山が高速検索の候補リストに絞り込まれ、次に再ランキングが
その候補リストを並べ替えて、正しい答えが
先頭へ移動する

この cookbook はその構成を、裁判所意見のデータセットで検証します。以下に示す再ランキングの例で行います。

高速検索とは何か

高速検索とは、大きなコーパス内のすべての文書とクエリを比較し、順位付けされた候補リストを素早く返せる任意の方法です。一般的な方法には、BM25 のようなキーワード検索や、パッセージを意味で比較する密な埋め込みがあります。システムはしばしば両方の方法を組み合わせます。

ここでの最初のステップは BM25 のみで、ほかに何も使いません。BM25 は共有される単語でパッセージを順位付けします。このステップを単純に保つことで、cookbook の狙いである再ランキングに注意を向けられます。高速検索の方法の選択は本筋ではありません。再ランキングが見るのは、候補リストに入ったパッセージだけです。

再ランキングとは何か

再ランキングは、高速検索がすでに作った候補リストを受け取り、より良い順序に並べ替えます。クエリをコーパス全体と一度に比較するのではなく、クエリを候補リスト上の各候補と個別に比較し、そのスコアで候補リストを並べ替えます。

図：左に順位付けされた候補リスト、中央に「再ランキング」と書かれた矢印、右に並べ替え後の
版があり、正しい答えが中央から
先頭へ移動する

スコアは言語モデルから得ることもできます。クエリと 1 つの候補を一緒に与え、その候補がクエリにどれだけよく答えるかを尋ねます。すると再ランキングは、候補の言い回しがクエリと異なっていても、候補リスト上の最良の一致を見つけられます。

TypeSafe による再ランキング

再ランカーは、クエリと候補の各ペアに対して比較可能なスコアを必要とします。汎用の言語モデルでもこれらのスコアを出せますし、候補リスト全体を直接順位付けすることもできます。しかしペアごとの独立したスコアリングでは、スコアの尺度を定義し、すべての候補に同じ基準を適用するようモデルに促す必要があります。同じペアでも呼び出しを繰り返すと異なるスコアが出ることがあり、汎用の生成は、1 つの数値だけを必要とするタスクに時間とコストを加えてしまいます。

TypeSafe が返すもの

TypeSafe では、スコアリングのリクエストをはい/いいえの質問のままにできます。

Could this candidate passage be from the cited precedent?

単純なはい/いいえでは 30 個の候補を順位付けするのに十分ではありません。代わりに Noul は 0 から 1 の間の数値を返します。これは noul と呼ばれます。noul は、答えが「はい」である確からしさを TypeSafe が推定したものです。

質問の criteria が何を真とし何を偽とするかを定義します。TypeSafe はそれをすべてのクエリ・候補ペアに適用し、noul を直接返します。その noul が、アプリケーションが並べ替えに使うスコアです。汎用モデル向けにスコアの尺度を発明する必要はなく、 TypeSafe はこの反復的なスコアリングをより速く、安く、一貫して行うように作られています。

簡略化した擬似コードでは、TypeSafe のスコアリング呼び出し 1 回は次のようになります。

question = Noul(
    instructions="Is this candidate the cited case?",
    criteria=NoulCriteria(
        true="The candidate states the specific rule the query cites.",
        false="The candidate is only on a similar topic.",
    ),
)
response = client.system_one(state={...}, questions={"is_cited_source": question})
response.answers["is_cited_source"].noul  # -> 0.87

TypeSafe はクエリと 1 つの候補をその質問に照らして一緒に読み、 noul を返します。

これを使って、候補リストの各候補に対して同じ質問を実行し、各呼び出しが返す noul で候補リストを並べ替え、最も高いものを先頭にすることで、候補リストを再ランキングできます。

nouls = {candidate: ask_typesafe(query, candidate) for candidate in shortlist}
reranked = sorted(shortlist, key=lambda c: nouls[c], reverse=True)  # highest noul first

以下の図は、候補ごとに 1 回のリクエストを行うことで、候補リストを並べ替えるためのスコアがどう生成されるかを示します。

flowchart LR
    q["query excerpt<br/><i>one opinion passage,<br/>citation removed</i>"]
    sl["shortlist from fast search<br/><i>30 candidate passages</i>"]
    quest["<b>one Noul</b><br/>could this candidate be<br/>from the cited precedent?<br/><i>criteria fix true and false</i>"]

    %% direction LR inside an LR chart keeps each state beside its noul, two columns,
    %% so the fan-out is four rows tall instead of eight
    subgraph fan["one request per candidate · no request sees another"]
        direction LR
        d1["state<br/>{query, candidate 1}"] --> n1["noul<br/>0.87"]
        d2["state<br/>{query, candidate 2}"] --> n2["noul<br/>0.41"]
        dx["⋮"] --> nx["⋮"]
        d30["state<br/>{query, candidate 30}"] --> n30["noul<br/>0.12"]
    end

    sort["sort by noul,<br/>highest first"]
    out["re-ranked shortlist<br/><i>same 30, better order</i>"]

    q --> fan
    sl --> fan
    quest --> fan
    fan --> sort --> out

    %% the elision is not a node - drop its box so it reads as "and so on"
    classDef elide fill:none,stroke:none
    class dx,nx elide
    linkStyle 2 stroke:none

再ランキングの例

次に、高速検索と再ランキングを、法律検索データセットである CLERC で実行します。この例では 3,565 件の裁判所意見のパッセージと 40 件のクエリを使います。

セットアップ

最初のステップは、このチュートリアルが依存するパッケージをインストールします。

bm25s と datasets が高速検索の候補リストを作ります。
typesafe-sdk と cooksafe が再ランキングと API キャッシュを担当します。
matplotlib が結果のグラフを描画します。

pip install bm25s datasets matplotlib 'cooksafe>=0.2.0,<0.3.0'

次のブロックは、TypeSafe クライアントと、ウォークスルーの残りが使う定数を設定します。たとえば、どの TypeSafe モデルを呼び出すか、高速検索が再ランカーに渡す候補リストの大きさなどです。TypeSafe の呼び出しには TYPESAFE_API_KEY が必要です。

import hashlib
import json
import os
import random
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

from cooksafe import JsonCache
from IPython.display import display
from typesafe_sdk import Noul, NoulCriteria, TypeSafeClient

TYPESAFE_MODEL = "jev-1.12"
PRICE = (
    0.042,
    0.00,
)  # $ per 1M tokens (input, output); TypeSafe jev-1.12 as of 2026-08
N_ROWS = 170  # CLERC rows pooled into the shared corpus
N_QUERIES = 40  # rows we evaluate
TOP_K = 30  # candidates the shortlist hands to the re-ranker, per query

client = TypeSafeClient(
    api_key=os.environ.get(
        "TYPESAFE_API_KEY", "cache-only"
    ),  # keyless kernels replay the cache
    base_url=os.environ.get("TYPESAFE_ENDPOINT"),
    timeout=120.0,
)
json_cache = JsonCache(Path("json_cache.json"))

高速検索でパッセージを順位付けする

ここで使うデータセットは米国の裁判所意見のコーパスで、170 行をまとめたものです。各行は次のように構成されます。

Query：引用を除いた意見の抜粋。
Gold：除かれた引用が指していたパッセージで、そのクエリに対する唯一の正解。
Candidates：コーパス内の他のすべてのパッセージで、それぞれクエリが誤ってマッチしうるもの。

170 行のうち 40 行をクエリとして評価するために選びます。残りの 130 行は候補としてだけ現れます。

次のセルは、上で説明した手法を使って候補リストを作ります。

コーパスを読み込みます。
BM25 でコーパスを全クエリに対して順位付けします。

ここではまだ TypeSafe は使いません。これは高速検索のステップだけです。

CLERC_FILE = (
    "https://huggingface.co/datasets/jhu-clsp/CLERC/resolve/main/"
    "teva_train_dir/train_data.jsonl.gz"
)

def cid(text: str) -> str:
    """Corpus id: a content hash, so passages shared across queries dedupe."""
    return hashlib.sha1(text.encode("utf-8")).hexdigest()[:16]

@json_cache
def build_slice(n_rows: int, n_queries: int, seed: int) -> dict:
    """Stream CLERC rows, pool ``n_rows`` of them into a corpus, pick ``n_queries`` to evaluate."""
    from datasets import load_dataset  # heavy import, keep local

    stream = load_dataset("json", data_files=CLERC_FILE, streaming=True, split="train")
    rows = []
    for row in stream:
        if (
            row.get("positive_passages")
            and len(row.get("negative_passages") or []) == 20
        ):
            rows.append(row)
        if len(rows) >= 1000:
            break

    rng = random.Random(seed)
    picked = rng.sample(rows, n_rows)
    corpus, pool = {}, []
    for row in picked:
        gold = row["positive_passages"][0]["text"]
        corpus[cid(gold)] = gold
        for neg in row["negative_passages"]:
            corpus[cid(neg["text"])] = neg["text"]
        pool.append(
            {"qid": str(row["query_id"]), "query": row["query"], "gold": cid(gold)}
        )
    # hold out the first 20 pooled rows; evaluate on the rest
    queries = rng.sample(pool[20:], n_queries)
    # sort the corpus by id so every run — live or cache replay — iterates it identically
    return {"queries": queries, "corpus": dict(sorted(corpus.items()))}

def bm25_rankings(corpus: dict[str, str], queries: dict[str, str], k: int = 100):
    """Rank every passage in the corpus by word overlap with each query."""
    import bm25s

    cids = list(corpus)
    retriever = bm25s.BM25()
    retriever.index(bm25s.tokenize([corpus[c] for c in cids], stopwords="en"))
    qids = list(queries)
    idxs, _ = retriever.retrieve(
        bm25s.tokenize([queries[q] for q in qids], stopwords="en"), k=min(k, len(cids))
    )
    return {q: [cids[i] for i in idxs[row]] for row, q in enumerate(qids)}

def gold_rank(ranked: list[str], gold: str) -> int | None:
    """1-based rank of the gold id, or None if it isn't in the list."""
    return ranked.index(gold) + 1 if gold in ranked else None

SURFACE, INK, INK2, MUTED = "#f8f8f2", "#34342f", "#34342f", "#7c7c77"
GRID, AXIS, BLUE, GREEN = "#d8d8cf", "#d8d8cf", "#5d76a2", "#6f9b52"

def bar_chart(labels: list[str], shares: list[float], title: str) -> None:
    """A small single-series bar chart of shares (0-1, shown as percentages)."""
    import matplotlib.pyplot as plt

    fig, ax = plt.subplots(figsize=(5, 3.2), facecolor=SURFACE)
    ax.set_facecolor(SURFACE)
    for side in ("top", "right"):
        ax.spines[side].set_visible(False)
    for side in ("left", "bottom"):
        ax.spines[side].set_color(AXIS)
    ax.tick_params(colors=MUTED, labelcolor=INK2, labelsize=9)
    ax.set_axisbelow(True)
    ax.grid(axis="y", color=GRID, linewidth=0.8)

    bars = ax.bar(labels, shares, width=0.55, color=[BLUE, GREEN][: len(labels)])
    ax.bar_label(
        bars,
        labels=[f"{s * 100:.0f}%" for s in shares],
        padding=4,
        color=INK,
        fontsize=11,
    )
    ax.set_ylim(0, 1.1)
    ax.set_yticks([0, 0.25, 0.5, 0.75, 1.0])
    ax.set_yticklabels(["0%", "25%", "50%", "75%", "100%"])
    ax.set_ylabel(f"share of {len(queries)} queries", color=INK2, fontsize=9)
    ax.set_title(title, loc="left", color=INK, fontsize=11)
    plt.tight_layout()
    display(fig)
    plt.close(fig)

ds = build_slice(N_ROWS, N_QUERIES, seed=0)
corpus: dict[str, str] = ds["corpus"]
queries = {q["qid"]: q["query"] for q in ds["queries"]}
golds = {q["qid"]: q["gold"] for q in ds["queries"]}

candidates = {q: ranked[:TOP_K] for q, ranked in bm25_rankings(corpus, queries).items()}

in_top_k = sum(golds[q] in candidates[q] for q in queries)
at_rank_1 = sum(candidates[q][0] == golds[q] for q in queries)

bar_chart(
    [f"In top {TOP_K}", "At rank 1"],
    [in_top_k / len(queries), at_rank_1 / len(queries)],
    f"Where the correct passage lands, {len(queries)} queries against {len(corpus):,} candidates",
)

高速検索が正しいパッセージを 1 位にするのはまれ

このグラフは、3,565 個の候補のうち、高速検索が正しいパッセージをどこに置くかを示します。

高速検索は、正解を含む候補リストにコーパスを確実に絞り込みます。 40 件のクエリすべてで正解を含みます。しかしそのパッセージが候補リストの最上位になることはまれで、わずか 5% です。

以下の再ランキングは、すでに候補リストにある上位 30 候補を並べ替えるだけです。高速検索が選ばなかったパッセージを加えることはできません。ここでは、候補リストが 40 件のクエリすべてで正しいパッセージを含むので、再ランキングは各パッセージをより良い位置に置くことに集中できます。

TypeSafe で再ランキングする

再ランキングは、候補リスト上の各候補をそのクエリに対してスコアリングし、そのスコアで並べ替えます。TypeSafe が各ペアについて尋ねる質問は、その候補が、クエリから除かれた引用が指すパッセージでありうるかどうかです。

次のセルは以下を行います。

その質問を定義します。
各候補リストで候補ごとに 1 回尋ねます。40 クエリ × 30 候補で合計 1,200 回の呼び出しを、順番にではなく並行に実行します。
各候補リストを、TypeSafe が返すスコアで並べ替え、再ランキング結果を生成します。

is_cited_source = Noul(
    instructions=(
        "The query excerpt comes from a US federal court opinion and was written "
        "immediately around a citation to a precedent; the citation itself has been "
        "removed. Could the candidate passage be from that cited precedent — does it "
        "establish the specific legal proposition the query excerpt invokes at its "
        "citation point?"
    ),
    criteria=NoulCriteria(
        true=(
            "The candidate passage states or establishes the specific rule, standard, "
            "holding, or fact pattern that the query excerpt attributes to its removed "
            "citation."
        ),
        false=(
            "The candidate passage is merely on a similar topic or doctrine; it does not "
            "supply the specific proposition the query excerpt relies on."
        ),
    ),
)

@json_cache
def score_candidate(model: str, query: str, candidate: str, question_json: str) -> dict:
    """One TypeSafe call about one (query, candidate) pair: a noul, plus token usage."""
    # the SDK takes a question as its JSON dict, so the cached string decodes straight in
    question = json.loads(question_json)
    response = client.system_one(
        state={"query_excerpt": query, "candidate_passage": candidate},
        questions={"is_cited_source": question},
        model=model,
    )
    return {
        "noul": response.answers["is_cited_source"].noul,
        "input_tokens": response.usage.input_tokens or 0,
        "output_tokens": response.usage.output_tokens or 0,
    }

# Each of the 40 queries has 30 candidates, so re-ranking every shortlist means 1,200 independent
# calls — cheap enough to fire all at once with a thread pool instead of one after another.
pair_list = [(q, c) for q in queries for c in candidates[q]]
question_json = is_cited_source.model_dump_json(exclude_none=True)
with ThreadPoolExecutor(max_workers=12) as pool:
    results = pool.map(
        lambda p: score_candidate(
            TYPESAFE_MODEL, queries[p[0]], corpus[p[1]], question_json
        ),
        pair_list,
    )
pair_scores = {q: {} for q in queries}
for (q, c), result in zip(pair_list, results):
    pair_scores[q][c] = result

reranked = {
    q: sorted(candidates[q], key=lambda c: -pair_scores[q][c]["noul"]) for q in queries
}

def chart_before_after(
    runs: dict[str, dict[str, list[str]]], thresholds: list[int]
) -> None:
    """Grouped bar chart: how often the correct passage lands in the top N, for each run."""
    import numpy as np
    import matplotlib.pyplot as plt

    labels = list(runs)
    colors = [BLUE, GREEN]

    def share_in_top(rankings, k):
        return sum(
            gold_rank(rankings[q], golds[q]) in range(1, k + 1) for q in queries
        ) / len(queries)

    fig, ax = plt.subplots(figsize=(6.5, 3.6), facecolor=SURFACE)
    ax.set_facecolor(SURFACE)
    for side in ("top", "right"):
        ax.spines[side].set_visible(False)
    for side in ("left", "bottom"):
        ax.spines[side].set_color(AXIS)
    ax.tick_params(colors=MUTED, labelcolor=INK2, labelsize=9)
    ax.set_axisbelow(True)
    ax.grid(axis="y", color=GRID, linewidth=0.8)

    x = np.arange(len(thresholds))
    width = 0.35
    for i, (label, rankings) in enumerate(runs.items()):
        shares = [share_in_top(rankings, k) for k in thresholds]
        offset = (i - (len(labels) - 1) / 2) * width
        bars = ax.bar(x + offset, shares, width * 0.92, color=colors[i], label=label)
        ax.bar_label(
            bars,
            labels=[f"{s * 100:.0f}%" for s in shares],
            padding=3,
            color=INK2,
            fontsize=8.5,
        )

    ax.set_xticks(x, [f"top {k}" for k in thresholds])
    ax.set_ylim(0, 1)
    ax.set_yticks([0, 0.25, 0.5, 0.75, 1.0])
    ax.set_yticklabels(["0%", "25%", "50%", "75%", "100%"])
    ax.set_ylabel(f"share of {len(queries)} queries", color=INK2, fontsize=9)
    ax.set_title(
        "How often the correct passage lands near the top",
        loc="left",
        color=INK,
        fontsize=11,
    )
    ax.legend(frameon=False, labelcolor=INK2, fontsize=9, loc="upper left")
    plt.tight_layout()
    display(fig)
    plt.close(fig)

chart_before_after(
    {"Fast search": candidates, "+ TypeSafe re-rank": reranked}, [1, 5, 10]
)

calls = [pair_scores[q][c] for q in queries for c in pair_scores[q]]
input_tokens = sum(call["input_tokens"] for call in calls)
output_tokens = sum(call["output_tokens"] for call in calls)
cost = input_tokens / 1_000_000 * PRICE[0] + output_tokens / 1_000_000 * PRICE[1]
print(
    f"{len(calls)} TypeSafe calls used {input_tokens:,} input and "
    f"{output_tokens:,} output tokens, costing ${cost:.4f}."
)

1200 TypeSafe calls used 1,536,002 input and 25,200 output tokens, costing $0.0645.

再ランキングが正しい答えを上へ引き上げる

このグラフは、高速検索と、高速検索＋再ランキングを 3 つのしきい値で比較します。再ランキングは、そのすべてで正しいパッセージを先頭に近づけます。

Top 1 — 5% → 18%
Top 5 — 15% → 35%
Top 10 — 38% → 62%

報告したトークン数とコストは、40 件の候補リストを再ランキングするために使った 1,200 回の TypeSafe 呼び出しすべてを対象としています。

各 CLERC 行は、1 つの正解パッセージと 20 の負例パッセージを含みます。このウォークスルーは 170 行のパッセージを 1 つの共有コーパスにまとめます。40 件の評価クエリそれぞれについて、 BM25 はその行に付属する 20 の負例だけでなく、完全なコーパスから 30 個の候補を選びます。次に TypeSafe がクエリを選択された各候補と照らして読み、その 30 パッセージを再ランキングします。

このウォークスルーでは、分かりやすさのために 1 ペアにつき 1 問だけ尋ねました。実際のアプリケーションなら、同じペアについて複数の質問を 1 回の呼び出しで尋ねるでしょう。質問の並列実行の cookbook と、投機的ファンアウトのパターンを参照してください。

次のステップ

同じ構成要素は TypeSafe のドキュメントの他の場所にも登場します。

Noul：TypeSafe がはい/いいえの質問をスコアに変える仕組み。
投機的ファンアウト：1 つの文書について複数の質問を 1 回の呼び出しで尋ねる方法。
行単位の検索：キーワードではなく意味でコーパスを検索するもう 1 つの方法。