Passer au contenu principal
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 Retrieval-Augmented Generation (RAG). Vous récupérez un modèle de chat affiné depuis le W&B Models Registry, le remplacez dans un RagModel existant suivi dans Weave, évaluez l’application mise à jour avec weave.Evaluation, puis publiez le nouveau modèle RAG dans le Registre. Ce flux de travail s’adresse aux équipes qui entraînent et affinent des modèles avec W&B Models et souhaitent les intégrer dans des applications LLM suivies et évaluées avec Weave.
Il s’agit d’un notebook interactif. Vous pouvez l’exécuter en local ou utiliser les liens suivants :

Prérequis

Commencez par installer les bibliothèques requises, configurer vos clés API, vous connecter à W&B et créer un projet W&B.
  1. Installez weave, pandas, unsloth, wandb, litellm, pydantic, torch et faiss-gpu à l’aide de pip.
%%capture
!pip install weave wandb pandas pydantic litellm faiss-gpu
python
%%capture
!pip install unsloth
# Obtenez é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"
  1. 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 and Weave
os.environ["OPENAI_API_KEY"] = userdata.get(
    "OPENAI_API_KEY"
)  # OpenAI - pour les embeddings de récupération
os.environ["GEMINI_API_KEY"] = userdata.get(
    "GEMINI_API_KEY"
)  # Gemini - pour le modèle de chat de base
  1. Connectez-vous à W&B, puis 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 Models registre et implémenter UnslothLoRAChatModel

Dans ce scénario, la Model Team a déjà affiné le modèle Llama-3.2 à l’aide de la bibliothèque unsloth pour optimiser les performances, et le modèle est disponible dans le W&B Models Registry. Dans cette étape, vous récupérez le ChatModel affiné depuis le registre et le convertissez en weave.Model pour le rendre compatible avec le RagModel.
Le RagModel mentionné dans le code suivant 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 nécessaire pour télécharger un artifact depuis le W&B registre. ChatModel peut être remplacé de manière modulaire pour prendre en charge n’importe quel autre type de modèle de chat LLM au sein du RagModel. Pour plus d’informations, voir le modèle dans Weave.
Pour charger le ChatModel, utilisez unsloth.FastLanguageModel ou peft.AutoPeftModelForCausalLM avec adaptateurs, ce qui permet une intégration efficace dans l’application. Après avoir téléchargé le modèle depuis le registre, configurez la logique d’initialisation et de prédiction à l’aide de la méthode model_post_init. Le code requis pour cette étape est disponible dans l’onglet Use du registre, et vous pouvez le copier directement dans votre implémentation. Le code suivant définit la classe UnslothLoRAChatModel pour gérer, initialiser et utiliser le modèle Llama-3.2 affiné récupéré depuis le W&B Models Registry. UnslothLoRAChatModel utilise unsloth.FastLanguageModel pour une inférence optimisée. 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 des réponses. Pour adapter le code à votre cas d’usage, mettez à jour MODEL_REG_URL avec le bon chemin du registre pour votre modèle affiné et ajustez des paramètres comme max_seq_length ou dtype en fonction de votre matériel ou de vos requirements.
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):
        # nous pouvons simplement coller ceci 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 l'inférence native 2x 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  # Au choix ! Le RoPE Scaling est pris en charge automatiquement en interne !
dtype = (
    None  # None pour la détection automatique. Float16 pour Tesla T4, V100, Bfloat16 pour Ampere+
)
load_in_4bit = True  # Utiliser la quantification 4 bits pour réduire l'utilisation de la mémoire. Peut être 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 sans avoir à reconstruire l’ensemble du pipeline. Dans cette étape, vous récupérez le RagModel existant dans votre projet Weave et mettez à jour son ChatModel pour utiliser le modèle affiné. Remplacer le modèle de chat par le nouveau laisse les autres composants, comme la base de données vectorielle et les prompts, inchangés, ce qui préserve la structure globale de l’application tout en améliorant les performances. Le code suivant récupère l’objet RagModel à l’aide d’une référence du projet Weave. L’attribut chat_model du RagModel est ensuite mis à jour pour utiliser la nouvelle instance UnslothLoRAChatModel créée à l’étape précédente. Ensuite, le RagModel mis à jour est publié afin de créer une nouvelle version. Enfin, le code 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 (pas besoin 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 afin de référencer la nouvelle version lors de la prédiction
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 à confirmer que le nouveau modèle de chat affiné se comporte comme prévu dans l’application. À cette étape, vous évaluez les performances du RagModel mis à jour à l’aide d’une weave.Evaluation existante. Ce processus garantit que le nouveau modèle de chat affiné se comporte comme prévu dans l’application RAG. Pour simplifier l’intégration et permettre la collaboration entre les équipes Models et Apps, vous enregistrez les résultats de l’évaluation à la fois dans le run W&B du modèle et dans le workspace Weave. Dans Models :
  • Le résumé de l’évaluation est enregistré dans le run W&B utilisé pour télécharger le modèle de chat affiné. Cela inclut des métriques de synthèse et des graphiques affichés dans une vue Workspace pour l’analyse.
  • L’ID de trace de l’évaluation est ajouté à la configuration du run, avec un lien direct vers la page Weave afin d’améliorer la traçabilité pour la 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é comme colonne supplémentaire dans les traces d’évaluation pour fournir davantage de contexte.
Le code suivant montre comment récupérer un objet d’évaluation, exécuter l’évaluation avec le RagModel mis à jour, puis enregistrer les résultats à la fois dans W&B et dans Weave. Assurez-vous que la référence de l’évaluation (WEAVE_EVAL) correspond à la configuration de votre projet.
# MAGIC: nous pouvons simplement obtenir une évaluation avec un dataset d'évaluation et des évaluateurs, puis les 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
# enregistrer dans 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()

Enregistrer le nouveau modèle RAG dans le registre

Maintenant que vous avez évalué le RagModel mis à jour, l’étape finale consiste à le publier dans le W&B Models Registry afin que d’autres équipes puissent le découvrir et le réutiliser. Pour rendre le RagModel mis à jour disponible pour une utilisation ultérieure par les équipes Models et Apps, nous le publions dans le W&B Models Registry en tant qu’artifact de référence. Le code suivant récupère la version et le nom de l’objet weave du RagModel mis à jour, puis les utilise pour créer des liens de référence. Le code crée ensuite un nouvel artifact dans W&B avec des métadonnées contenant l’URL Weave du modèle. Le code enregistre cet artifact dans le registre W&B et le lie à un chemin de registre spécifié. Avant d’exécuter le code, assurez-vous que les variables ENTITY et PROJECT correspondent à votre configuration W&B, et spécifiez le chemin de registre cible correct. Ce processus finalise le flux de travail en publiant le nouveau RagModel dans l’écosystème W&B pour la collaboration et la réutilisation. Après avoir exécuté le code de cette section, votre RagModel mis à jour est disponible dans le W&B Models Registry en tant qu’artifact référencé, ce qui achève 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)

    # enregistrer 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"
    )