メインコンテンツへスキップ
これはインタラクティブなノートブックです。ローカルで実行するか、以下のリンクを使用できます。

​
Not Diamond を使った LLM プロンプトのカスタムルーティング

このノートブックでは、Weave と Not Diamond のカスタムルーティング を使って、評価結果に基づき、LLM プロンプトを最適なモデルにルーティングする方法を説明します。

​
プロンプトのルーティング

複雑な LLM ワークフローを構築する際、ユーザーは精度、コスト、または call レイテンシに応じて、異なるモデルにプロンプトを振り分ける必要がある場合があります。 Users は Not Diamond を使用して、こうしたワークフロー内のプロンプトをニーズに合った適切なモデルにルーティングできるため、モデルのコストを抑えつつ精度を最大化できます。 どのようなデータ分布であっても、1 つのモデルがすべてのクエリで他のすべてのモデルを常に上回ることは、ほとんどありません。複数のモデルを組み合わせて、各 LLM をいつ call するかを学習する「メタモデル」を作ることで、個々のモデルそれぞれを上回るパフォーマンスを実現し、さらにコストとレイテンシを削減することもできます。

​
カスタムルーティング

プロンプト用のカスタムルーターをトレーニングするには、3 つのものが必要です。
  1. LLM プロンプトのセット: プロンプトは文字列である必要があり、アプリケーションで使用するプロンプトを適切に表している必要があります。
  2. LLM 応答: 各入力に対する候補 LLM からの応答です。候補 LLM には、サポート対象の LLM と独自のカスタムモデルの両方を含めることができます。
  3. 候補 LLM による各入力への応答の評価スコア: スコアは数値で、ニーズに合った任意のメトリクスを使用できます。
これらを Not Diamond API に送信すれば、各ワークフローに合わせて調整したカスタムルーターをトレーニングできます。

​
トレーニングデータのセットアップ

通常は、独自の評価を使用してカスタムルーターをトレーニングします。ただし、このサンプルノートブックでは、コーディングタスク向けのカスタムルーターをトレーニングするために、HumanEvalデータセット に対する LLM の応答を使用します。 まず、この例のために用意したデータセットをダウンロードし、次に LLM の応答を各モデルの EvaluationResults にパースします。 
!curl -L "https://drive.google.com/uc?export=download&id=1q1zNZHioy9B7M-WRjsJPkfvFosfaHX38" -o humaneval.csv
python
import random

import weave
from weave.flow.dataset import Dataset
from weave.flow.eval import EvaluationResults
from weave.integrations.notdiamond.util import get_model_evals

pct_train = 0.8
pct_test = 1 - pct_train

# 実際には、データセットに対して評価を構築し、
# `evaluation.get_eval_results(model)` を呼び出します
model_evals = get_model_evals("./humaneval.csv")
model_train = {}
model_test = {}
for model, evaluation_results in model_evals.items():
    n_results = len(evaluation_results.rows)
    all_idxs = list(range(n_results))
    train_idxs = random.sample(all_idxs, k=int(n_results * pct_train))
    test_idxs = [idx for idx in all_idxs if idx not in train_idxs]

    model_train[model] = EvaluationResults(
        rows=weave.Table([evaluation_results.rows[idx] for idx in train_idxs])
    )
    model_test[model] = Dataset(
        rows=weave.Table([evaluation_results.rows[idx] for idx in test_idxs])
    )
    print(
        f"Found {len(train_idxs)} train rows and {len(test_idxs)} test rows for {model}."
    )

​
カスタムルーターをトレーニングする

EvaluationResults が揃ったら、カスタムルーターをトレーニングできます。まず、アカウントを作成 し、 APIキーを生成してから、以下に APIキーを入力してください。
APIキーを作成
import os

from weave.integrations.notdiamond.custom_router import train_router

api_key = os.getenv("NOTDIAMOND_API_KEY", "<YOUR_API_KEY>")

preference_id = train_router(
    model_evals=model_train,
    prompt_column="prompt",
    response_column="actual",
    language="en",
    maximize=True,
    api_key=api_key,
    # 初めてカスタムルーターをトレーニングする場合はこのままにしてください
    # カスタムルーターを再トレーニングする場合はコメントを解除してください
    # preference_id=preference_id,
)
その後、Not Diamond アプリからカスタムルーターのトレーニングの進行状況を確認できます。
ルーターのトレーニング進捗の確認
カスタムルーターのトレーニングが完了したら、それを使用してプロンプトをルーティングできます。 
from notdiamond import NotDiamond

import weave

weave.init("notdiamond-quickstart")

llm_configs = [
    "anthropic/claude-3-5-sonnet-20240620",
    "openai/gpt-4o-2024-05-13",
    "google/gemini-1.5-pro-latest",
    "openai/gpt-4-turbo-2024-04-09",
    "anthropic/claude-3-opus-20240229",
]
client = NotDiamond(api_key=api_key, llm_configs=llm_configs)

new_prompt = (
    """
You are a helpful coding assistant. Using the provided function signature, write the implementation for the function
in Python. Write only the function. Do not include any other text.

from typing import List

def has_close_elements(numbers: List[float], threshold: float) -> bool:
    """
    """ Check if in given list of numbers, are any two numbers closer to each other than
    given threshold.
    >>> has_close_elements([1.0, 2.0, 3.0], 0.5)
    False
    >>> has_close_elements([1.0, 2.8, 3.0, 4.0, 5.0, 2.0], 0.3)
    True
    """
    """
"""
)
session_id, routing_target_model = client.model_select(
    messages=[{"role": "user", "content": new_prompt}],
    preference_id=preference_id,
)

print(f"Session ID: {session_id}")
print(f"Target Model: {routing_target_model}")
この例では、Not Diamond の Weave 自動トレースとの互換性も活用しています。結果は Weave UI で確認できます。 カスタムルーティングの Weave UI

​
カスタムルーターの評価

カスタムルーターをトレーニングしたら、次のいずれかのパフォーマンスを評価できます。
  • トレーニング用プロンプトを送信して、インサンプルのパフォーマンスを評価する
  • 新しいプロンプトまたはホールドアウトしたプロンプトを送信して、アウトオブサンプルのパフォーマンスを評価する
以下では、テストセットをカスタムルーターに送信して、そのパフォーマンスを評価します。 
from weave.integrations.notdiamond.custom_router import evaluate_router

eval_prompt_column = "prompt"
eval_response_column = "actual"

best_provider_model, nd_model = evaluate_router(
    model_datasets=model_test,
    prompt_column=eval_prompt_column,
    response_column=eval_response_column,
    api_key=api_key,
    preference_id=preference_id,
)
python
@weave.op()
def is_correct(score: int, output: dict) -> dict:
    # すでにモデルの応答があるため、スコアをそのまま使用する
    return {"correct": score}

best_provider_eval = weave.Evaluation(
    dataset=best_provider_model.model_results.to_dict(orient="records"),
    scorers=[is_correct],
)
await best_provider_eval.evaluate(best_provider_model)

nd_eval = weave.Evaluation(
    dataset=nd_model.model_results.to_dict(orient="records"), scorers=[is_correct]
)
await nd_eval.evaluate(nd_model)
この例では、Not Diamond の “メタモデル” が、複数の異なるモデルにプロンプトを振り分けます。 Weave 経由でカスタムルーターをトレーニングすると、評価も実行され、結果が Weave UI にアップロードされます。カスタムルーターの処理が完了すると、Weave UI で結果を確認できます。 UI では、Not Diamond の “メタモデル” が、プロンプトに正確に回答できる可能性が高い別のモデルへプロンプトを振り分けることで、最高性能のモデルを上回っていることがわかります。
Not Diamond の評価