Models and Weave Integration Demo

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

• Google Colabで開く
• GitHubでソースを表示

​

前提条件

まず、必要なライブラリをインストールし、APIキーを設定し、W&Bにログインして、新しいW&Bプロジェクトを作成します。

1. インストール weave、 pandas、 unsloth、 wandb、 litellm、 pydantic、 torch、および faiss-gpu を使用してpip。

%%capture
!pip install weave wandb pandas pydantic litellm faiss-gpu
python
%%capture
!pip install unsloth
# Also get the latest nightly Unsloth!
!pip uninstall unsloth -y && pip install --upgrade --no-cache-dir "unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git"

2. 環境から必要なAPIキーを追加します。

import os

from google.colab import userdata

os.environ["WANDB_API_KEY"] = userdata.get("WANDB_API_KEY")  # W&B Models and Weave
os.environ["OPENAI_API_KEY"] = userdata.get(
    "OPENAI_API_KEY"
)  # OpenAI - for retrieval embeddings
os.environ["GEMINI_API_KEY"] = userdata.get(
    "GEMINI_API_KEY"
)  # Gemini - for the base chat model

3. W&Bにログインし、新しいプロジェクトを作成します。

import pandas as pd
import wandb

import weave

wandb.login()

PROJECT = "weave-cookboook-demo"
ENTITY = "wandb-smle"

weave.init(ENTITY + "/" + PROJECT)

​

ダウンロード ChatModel をModels Registryから実装 UnslothLoRAChatModel

このシナリオでは、Llama-3.2モデルはすでにModel Teamによって unsloth ライブラリを使用してパフォーマンス最適化のためにファインチューニングされており、W&B Models Registryで利用可能です。このステップでは、Registryからファインチューニングされた ChatModel を取得し、weave.Model に変換して RagModel と互換性を持たせます。

以下で参照される RagModel は、完全なRAGアプリケーションと見なすことができるトップレベルの weave.Model です。これには ChatModel、ベクトルデータベース、およびプロンプトが含まれています。ChatModel も weave.Model であり、W&B Registryからアーティファクトをダウンロードするコードが含まれています。ChatModel は、RagModel の一部として他の種類のLLMチャットモデルをサポートするようにモジュラーに変更できます。詳細については、Weaveでモデルを表示してください。

ロードするには ChatModel、unsloth.FastLanguageModel または peft.AutoPeftModelForCausalLM とアダプターが使用され、アプリへの効率的な統合が可能になります。Registryからモデルをダウンロードした後、model_post_init メソッドを使用して初期化と予測ロジックを設定できます。このステップに必要なコードは Use タブのRegistryで利用可能で、実装に直接コピーできます

以下のコードは UnslothLoRAChatModel クラスを定義して、W&B Models Registryから取得したファインチューニングされたLlama-3.2モデルを管理、初期化、使用します。UnslothLoRAChatModel は unsloth.FastLanguageModel を最適化された推論に使用します。model_post_init メソッドはモデルのダウンロードとセットアップを処理し、predict メソッドはユーザークエリを処理し、応答を生成します。ユースケースに合わせてコードを適応させるには、MODEL_REG_URL をファインチューニングされたモデルの正しいRegistryパスで更新し、max_seq_length や dtype などのパラメータをハードウェアや要件に基づいて調整します。

from typing import Any

from pydantic import PrivateAttr
from unsloth import FastLanguageModel

import weave

class UnslothLoRAChatModel(weave.Model):
    """
    We define an extra ChatModel class to be able store and version more parameters than just the model name.
    Especially, relevant if we consider fine-tuning (locally or aaS) because of specific parameters.
    """

    chat_model: str
    cm_temperature: float
    cm_max_new_tokens: int
    cm_quantize: bool
    inference_batch_size: int
    dtype: Any
    device: str
    _model: Any = PrivateAttr()
    _tokenizer: Any = PrivateAttr()

    def model_post_init(self, __context):
        # we can simply paste this from the "Use" tab from the registry
        run = wandb.init(project=PROJECT, job_type="model_download")
        artifact = run.use_artifact(f"{self.chat_model}")
        model_path = artifact.download()

        # unsloth version (enable native 2x faster inference)
        self._model, self._tokenizer = FastLanguageModel.from_pretrained(
            model_name=model_path,
            max_seq_length=self.cm_max_new_tokens,
            dtype=self.dtype,
            load_in_4bit=self.cm_quantize,
        )
        FastLanguageModel.for_inference(self._model)

    @weave.op()
    async def predict(self, query: list[str]) -> dict:
        # add_generation_prompt = true - Must add for generation
        input_ids = self._tokenizer.apply_chat_template(
            query,
            tokenize=True,
            add_generation_prompt=True,
            return_tensors="pt",
        ).to("cuda")

        output_ids = self._model.generate(
            input_ids=input_ids,
            max_new_tokens=64,
            use_cache=True,
            temperature=1.5,
            min_p=0.1,
        )

        decoded_outputs = self._tokenizer.batch_decode(
            output_ids[0][input_ids.shape[1] :], skip_special_tokens=True
        )

        return "".join(decoded_outputs).strip()
python
MODEL_REG_URL = "wandb32/wandb-registry-RAG Chat Models/Finetuned Llama-3.2:v3"

max_seq_length = 2048  # Choose any! We auto support RoPE Scaling internally!
dtype = (
    None  # None for auto detection. Float16 for Tesla T4, V100, Bfloat16 for Ampere+
)
load_in_4bit = True  # Use 4bit quantization to reduce memory usage. Can be False.

new_chat_model = UnslothLoRAChatModel(
    name="UnslothLoRAChatModelRag",
    chat_model=MODEL_REG_URL,
    cm_temperature=1.0,
    cm_max_new_tokens=max_seq_length,
    cm_quantize=load_in_4bit,
    inference_batch_size=max_seq_length,
    dtype=dtype,
    device="auto",
)
python
await new_chat_model.predict(
    [{"role": "user", "content": "What is the capital of Germany?"}]
)

​

新しい ChatModel バージョンを RagModel

に統合するRagModel をWeaveプロジェクトから取得し、その ChatModel を更新して新しくファインチューニングされたモデルを使用します。このシームレスな交換により、ベクトルデータベース（VDB）やプロンプトなどの他のコンポーネントは変更されず、アプリケーションの全体的な構造を維持しながらパフォーマンスを向上させることができます。

以下のコードは、Weaveプロジェクトからの参照を使用して RagModel オブジェクトを取得します。chat_model 属性の RagModel は、前のステップで作成された新しい UnslothLoRAChatModel インスタンスを使用するように更新されます。その後、更新された RagModel が公開され、新しいバージョンが作成されます。最後に、更新された RagModel を使用してサンプル予測クエリを実行し、新しいチャットモデルが使用されていることを確認します。

RagModel = weave.ref(
    "weave:///wandb-smle/weave-cookboook-demo/object/RagModel:cqRaGKcxutBWXyM0fCGTR1Yk2mISLsNari4wlGTwERo"
).get()
python
RagModel.chat_model.chat_model
python
await RagModel.predict("When was the first conference on climate change?")
python
# MAGIC: exchange chat_model and publish new version (no need to worry about other RAG components)
RagModel.chat_model = new_chat_model
python
RagModel.chat_model.chat_model
python
# first publish new version so that in prediction we reference new version
PUB_REFERENCE = weave.publish(RagModel, "RagModel")
python
await RagModel.predict("When was the first conference on climate change?")

​

実行 weave.Evaluation

次のステップでは、既存の RagModel を使用して更新された weave.Evaluation のパフォーマンスを評価します。このプロセスにより、新しくファインチューニングされたチャットモデルがRAGアプリケーション内で期待通りに機能していることを確認します。ModelsチームとAppsチーム間の統合とコラボレーションを効率化するために、評価結果をモデルのW&B実行の一部としてもWeaveワークスペースの一部としても記録します。

Modelsでは：

• 評価の概要は、ファインチューニングされたチャットモデルをダウンロードするために使用されたW&B実行に記録されます。これには、分析のための workspace view に表示される要約メトリクスとグラフが含まれます。
• 評価トレースIDが実行の構成に追加され、Modelチームによるトレーサビリティを容易にするためにWeaveページに直接リンクします。

Weaveでは：

• アーティファクトまたはレジストリリンク ChatModel は RagModel への入力として保存されます。
• W&B実行IDは、より良いコンテキストのために評価トレースの追加列として保存されます。

以下のコードは、評価オブジェクトを取得し、更新された RagModel を使用して評価を実行し、結果をW&BとWeaveの両方に記録する方法を示しています。評価参照（WEAVE_EVAL）がプロジェクト設定と一致していることを確認してください。

# MAGIC: we can simply get an evaluation with a eval dataset and scorers and use them
WEAVE_EVAL = "weave:///wandb-smle/weave-cookboook-demo/object/climate_rag_eval:ntRX6qn3Tx6w3UEVZXdhIh1BWGh7uXcQpOQnIuvnSgo"
climate_rag_eval = weave.ref(WEAVE_EVAL).get()
python
with weave.attributes({"wandb-run-id": wandb.run.id}):
    # use .call attribute to retrieve both the result and the call in order to save eval trace to Models
    summary, call = await climate_rag_eval.evaluate.call(climate_rag_eval, RagModel)
python
# log to models
wandb.run.log(pd.json_normalize(summary, sep="/").to_dict(orient="records")[0])
wandb.run.config.update(
    {"weave_url": f"https://wandb.ai/wandb-smle/weave-cookboook-demo/r/call/{call.id}"}
)
wandb.run.finish()

​

新しいRAGモデルをRegistryに保存する

更新された RagModel をModelsチームとAppsチームの両方が将来使用できるようにするために、参照アーティファクトとしてW&B Models Registryにプッシュします。

以下のコードは、更新された weave オブジェクトのバージョンと名前を取得し、RagModel を使用して参照リンクを作成します。次に、モデルのWeave URLを含むメタデータを持つ新しいアーティファクトがW&Bで作成されます。このアーティファクトはW&B Registryに記録され、指定されたレジストリパスにリンクされます。

コードを実行する前に、ENTITY と PROJECT 変数がW&B設定と一致し、ターゲットレジストリパスが正しく指定されていることを確認してください。このプロセスにより、新しい RagModel をW&Bエコシステムに公開し、簡単なコラボレーションと再利用を可能にするワークフローが完了します。

MODELS_OBJECT_VERSION = PUB_REFERENCE.digest  # weave object version
MODELS_OBJECT_NAME = PUB_REFERENCE.name  # weave object name
python
models_url = f"https://wandb.ai/{ENTITY}/{PROJECT}/weave/objects/{MODELS_OBJECT_NAME}/versions/{MODELS_OBJECT_VERSION}"
models_link = (
    f"weave:///{ENTITY}/{PROJECT}/object/{MODELS_OBJECT_NAME}:{MODELS_OBJECT_VERSION}"
)

with wandb.init(project=PROJECT, entity=ENTITY) as run:
    # create new Artifact
    artifact_model = wandb.Artifact(
        name="RagModel",
        type="model",
        description="Models Link from RagModel in Weave",
        metadata={"url": models_url},
    )
    artifact_model.add_reference(models_link, name="model", checksum=False)

    # log new artifact
    run.log_artifact(artifact_model, aliases=[MODELS_OBJECT_VERSION])

    # link to registry
    run.link_artifact(
        artifact_model, target_path="wandb32/wandb-registry-RAG Models/RAG Model"
    )