Utiliser Weave avec W&B Models

Ce notebook présente un flux de travail de bout en bout qui combine W&B Models et Weave pour créer, évaluer et publier une application de génération augmentée par récupération (RAG). Vous récupérez un modèle de chat affiné depuis le registre, vous le substituez au modèle d’un RagModel existant suivi dans Weave, vous évaluez l’application mise à jour avec weave.Evaluation, puis vous publiez le nouveau modèle RAG dans le registre. Ce flux de travail s’adresse aux équipes qui entraînent des modèles et en effectuent le fine-tuning avec W&B Models, et qui souhaitent les intégrer à des applications LLM suivies et évaluées avec Weave.

Il s’agit d’un notebook interactif. Vous pouvez l’exécuter localement ou utiliser les liens suivants :

Prérequis

Commencez par installer les bibliothèques requises, configurer les clés API, vous connecter à Forge et créer un projet Weights & Biases.

Installez weave, pandas, unsloth, wandb, litellm, pydantic, torch et faiss-gpu avec pip.

%%capture
!pip install weave wandb pandas pydantic litellm faiss-gpu
python
%%capture
!pip install unsloth
# Installez également la dernière version nightly d’Unsloth !
!pip uninstall unsloth -y && pip install --upgrade --no-cache-dir "unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git"

Ajoutez les clés API requises depuis votre environnement.

import os

from google.colab import userdata

os.environ["WANDB_API_KEY"] = userdata.get("WANDB_API_KEY")  # W&B Models et Weave
os.environ["OPENAI_API_KEY"] = userdata.get(
    "OPENAI_API_KEY"
)  # OpenAI - pour les embeddings de recherche
os.environ["GEMINI_API_KEY"] = userdata.get(
    "GEMINI_API_KEY"
)  # Gemini - pour le modèle de chat de base

Connectez-vous à Forge et créez un nouveau projet.

import pandas as pd
import wandb

import weave

wandb.login()

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

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

Télécharger `ChatModel` depuis le registre et implémenter `UnslothLoRAChatModel`

Dans ce scénario, la Model Team a déjà affiné le modèle Llama-3.2 avec la bibliothèque unsloth afin d’optimiser les performances, et le modèle est disponible dans W&B Registry. Au cours de cette étape, vous récupérez le ChatModel affiné depuis le registre, puis vous le convertissez en weave.Model pour le rendre compatible avec le RagModel.

Le RagModel utilisé dans le code ci-dessous est un weave.Model de premier niveau qui peut être considéré comme une application RAG complète. Il contient un ChatModel, une base de données vectorielle et un prompt. Le ChatModel est lui aussi un weave.Model, qui contient le code permettant de télécharger un artifact depuis W&B Registry. Grâce à cette approche modulaire, vous pouvez remplacer le ChatModel par n’importe quel autre modèle de chat LLM au sein du RagModel. Pour plus d’informations, consultez le modèle dans Weave.

Pour charger le ChatModel, utilisez unsloth.FastLanguageModel ou peft.AutoPeftModelForCausalLM avec des adaptateurs, afin de l’intégrer efficacement dans l’application. Une fois le modèle téléchargé depuis le registre, configurez la logique d’initialisation et de prédiction dans la méthode model_post_init. Le code nécessaire à cette étape est disponible dans l’onglet Use du registre : vous pouvez le copier directement dans votre implémentation. Le code suivant définit la classe UnslothLoRAChatModel, qui permet de gérer, d’initialiser et d’utiliser le modèle Llama-3.2 affiné récupéré depuis le registre. UnslothLoRAChatModel s’appuie sur unsloth.FastLanguageModel pour optimiser l’inférence. La méthode model_post_init télécharge et configure le modèle, tandis que la méthode predict traite les requêtes des utilisateurs et génère les réponses. Pour adapter le code à votre cas d’usage, remplacez la valeur de MODEL_REG_URL par le chemin de registre de votre modèle affiné, et ajustez des paramètres tels que max_seq_length ou dtype en fonction de votre matériel ou de vos besoins.

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):
        # il suffit de coller ce code depuis l’onglet « Use » du registre
        run = wandb.init(project=PROJECT, job_type="model_download")
        artifact = run.use_artifact(f"{self.chat_model}")
        model_path = artifact.download()

        # version unsloth (active nativement une inférence 2 fois plus rapide)
        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 - Obligatoire pour la génération
        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  # Choisissez la valeur de votre choix ! Le RoPE Scaling est géré automatiquement en interne !
dtype = (
    None  # None pour la détection automatique. Float16 pour Tesla T4 et V100, Bfloat16 pour Ampere+
)
load_in_4bit = True  # Quantification 4 bits pour réduire l’utilisation de la mémoire. Peut valoir 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?"}]
)

Intégrer la nouvelle version de `ChatModel` dans `RagModel`

Construire une application RAG à partir d’un modèle de chat affiné vous permet de réutiliser des composants adaptés à vos besoins sans reconstruire l’ensemble du pipeline. Au cours de cette étape, vous récupérez le RagModel existant depuis votre projet Weave et mettez à jour son ChatModel afin qu’il utilise le modèle affiné. Remplacer le modèle de chat laisse intacts les autres composants, comme la base de données vectorielle et les prompts : la structure globale de l’application est ainsi préservée, tout en améliorant ses performances. Le code suivant récupère l’objet RagModel à l’aide d’une référence issue du projet Weave. Il met ensuite à jour l’attribut chat_model du RagModel afin qu’il utilise la nouvelle instance UnslothLoRAChatModel créée à l’étape précédente. Il publie alors le RagModel mis à jour pour en créer une nouvelle version. Enfin, il exécute un exemple de requête de prédiction avec le RagModel mis à jour afin de vérifier qu’il utilise bien le nouveau modèle de chat.

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 : remplacer chat_model et publier une nouvelle version (inutile de se soucier des autres composants RAG)
RagModel.chat_model = new_chat_model
python
RagModel.chat_model.chat_model
python
# publier d'abord la nouvelle version pour que la prédiction y fasse référence
PUB_REFERENCE = weave.publish(RagModel, "RagModel")
python
await RagModel.predict("When was the first conference on climate change?")

Exécuter une `weave.Evaluation`

Une fois le RagModel mis à jour publié, l’étape suivante consiste à vérifier que le nouveau modèle de chat affiné se comporte comme prévu au sein de l’application. Dans cette étape, vous évaluez les performances du RagModel mis à jour à l’aide d’une weave.Evaluation existante. Vous confirmez ainsi que le nouveau modèle de chat affiné fonctionne comme prévu dans l’application RAG. Pour simplifier l’intégration et faciliter la collaboration entre les équipes Models et Apps, vous journalisez les résultats de l’évaluation à la fois dans le run W&B du modèle et dans le workspace Weave. Dans Models :

La synthèse de l’évaluation est journalisée dans le run W&B utilisé pour télécharger le modèle de chat affiné. Elle comprend les métriques de synthèse et les graphiques affichés dans une vue de workspace à des fins d’analyse.
L’ID de la trace d’évaluation est ajouté à la configuration du run et renvoie directement à la page Weave, ce qui améliore la traçabilité pour l’équipe Model Team.

Dans Weave :

Le lien vers l’artifact ou le registre du ChatModel est stocké comme entrée du RagModel.
L’ID du run W&B est enregistré dans une colonne supplémentaire des traces d’évaluation afin d’offrir davantage de contexte.

Le code suivant montre comment récupérer un objet d’évaluation, exécuter l’évaluation avec le RagModel mis à jour et journaliser les résultats dans W&B et Weave. Assurez-vous que la référence d’évaluation (WEAVE_EVAL) correspond à la configuration de votre projet.

# MAGIC : il suffit de récupérer une évaluation comportant un jeu de données d'évaluation et des évaluateurs, puis de l'utiliser
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}):
    # utiliser l'attribut .call pour récupérer à la fois le résultat et l'appel, afin d'enregistrer la trace d'évaluation dans Models
    summary, call = await climate_rag_eval.evaluate.call(climate_rag_eval, RagModel)
python
# journaliser dans Models
wandb.run.log(pd.json_normalize(summary, sep="/").to_dict(orient="records")[0])
wandb.run.config.update(
    {"weave_url": f"https://forge.coreweave.com/wandb/wandb-smle/weave-cookboook-demo/r/call/{call.id}"}
)
wandb.run.finish()

Enregistrer le nouveau modèle RAG dans le registre

Maintenant que vous avez évalué le RagModel mis à jour, la dernière étape consiste à le publier dans le registre afin que d’autres équipes puissent le découvrir et le réutiliser. Pour que les équipes Models et Apps puissent utiliser ultérieurement le RagModel mis à jour, téléversez-le dans le registre en tant qu’artifact de référence. Le code suivant récupère la version de l’objet weave et le nom du RagModel mis à jour, puis s’en sert pour créer des liens de référence. Il crée ensuite dans W&B un nouvel artifact dont les métadonnées contiennent l’URL Weave du modèle. Enfin, il journalise cet artifact dans le registre et le lie à un chemin de registre donné. Avant d’exécuter le code, vérifiez que les variables ENTITY et PROJECT correspondent à votre configuration W&B et indiquez le bon chemin de registre cible. Cette opération conclut le flux de travail en publiant le nouveau RagModel dans l’écosystème W&B pour faciliter la collaboration et la réutilisation. Une fois le code de cette section exécuté, votre RagModel mis à jour est disponible dans le registre en tant qu’artifact de référence, ce qui boucle l’aller-retour entre W&B Models et Weave.

MODELS_OBJECT_VERSION = PUB_REFERENCE.digest  # version de l’objet weave
MODELS_OBJECT_NAME = PUB_REFERENCE.name  # nom de l’objet weave
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:
    # créer un nouvel 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)

    # journaliser le nouvel artifact
    run.log_artifact(artifact_model, aliases=[MODELS_OBJECT_VERSION])

    # lier au registre
    run.link_artifact(
        artifact_model, target_path="wandb32/wandb-registry-RAG Models/RAG Model"
    )

​Prérequis

​Télécharger ChatModel depuis le registre et implémenter UnslothLoRAChatModel

​Intégrer la nouvelle version de ChatModel dans RagModel

​Exécuter une weave.Evaluation

​Enregistrer le nouveau modèle RAG dans le registre

Prérequis

Télécharger `ChatModel` depuis le registre et implémenter `UnslothLoRAChatModel`

Intégrer la nouvelle version de `ChatModel` dans `RagModel`

Exécuter une `weave.Evaluation`

Enregistrer le nouveau modèle RAG dans le registre