Passer au contenu principal
Il s’agit d’un notebook interactif. Vous pouvez l’exécuter localement ou utiliser les liens ci-dessous :

Configurer un modèle de coût personnalisé

Ce guide vous montre comment définir dans Weave un modèle personnalisé qui indique sa propre utilisation des jetons, enregistrer un coût personnalisé par jeton pour ce modèle, puis récupérer les appels avec les informations de coût associées. Utilisez cette approche lorsque vous travaillez avec un modèle qui n’est pas couvert par les données de coût intégrées de Weave et que vous souhaitez suivre les dépenses en parallèle de vos traces. Weave calcule les coûts en fonction du nombre de jetons consommés et du modèle utilisé. Weave récupère cette utilisation et ce modèle dans la sortie, puis les associe à l’appel. Configurez un modèle personnalisé simple qui calcule lui-même son utilisation des jetons et l’enregistre dans Weave.

Configurer l’environnement

Installez et importez tous les packages nécessaires. Définissez WANDB_API_KEY dans votre environnement afin de pouvoir vous connecter avec wandb.login() (cette clé doit être fournie à Colab en tant que secret). Définissez dans name_of_wandb_project le projet W&B dans lequel vous voulez journaliser ces données.
name_of_wandb_project peut également être au format {team_name}/{project_name} pour indiquer l’équipe dans laquelle journaliser les traces.
Récupérez ensuite un client Weave en appelant weave.init().
%pip install wandb weave datetime --quiet
python
import os

import wandb
from google.colab import userdata

import weave

os.environ["WANDB_API_KEY"] = userdata.get("WANDB_API_KEY")
name_of_wandb_project = "custom-cost-model"

wandb.login()
python
weave_client = weave.init(name_of_wandb_project)

Configurer un modèle avec Weave

Ensuite, définissez une sous-classe Weave Model qui effectue son propre comptage des jetons. Le fait de renvoyer les compteurs d’utilisation et le nom du modèle dans le dictionnaire de sortie permet ensuite à Weave d’associer l’appel au coût personnalisé défini dans la section suivante.
from weave import Model

class YourModel(Model):
    attribute1: str
    attribute2: int

    def simple_token_count(self, text: str) -> int:
        return len(text) // 3

    # Ceci est un op personnalisé que nous définissons
    # Il prend une chaîne en entrée et renvoie un dictionnaire avec les comptages d'utilisation, le nom du modèle et la sortie
    @weave.op()
    def custom_model_generate(self, input_data: str) -> dict:
        # La logique du modèle se trouve ici
        # C'est ici que vous placeriez une fonction de génération personnalisée
        prediction = self.attribute1 + " " + input_data

        # Comptages d'utilisation
        prompt_tokens = self.simple_token_count(input_data)
        completion_tokens = self.simple_token_count(prediction)

        # Nous renvoyons un dictionnaire avec les comptages d'utilisation, le nom du modèle et la sortie
        # Weave associera automatiquement ceci à la trace
        # Cet objet {usage, model, output} correspond à la sortie d'un appel OpenAI
        return {
            "usage": {
                "input_tokens": prompt_tokens,
                "output_tokens": completion_tokens,
                "total_tokens": prompt_tokens + completion_tokens,
            },
            "model": "your_model_name",
            "output": prediction,
        }

    # Dans notre fonction predict, nous appelons notre fonction de génération personnalisée et renvoyons la sortie.
    @weave.op()
    def predict(self, input_data: str) -> dict:
        # C'est ici que vous effectueriez tout post-traitement des données
        outputs = self.custom_model_generate(input_data)
        return outputs["output"]

Ajouter un coût personnalisé

Ajoutez un coût personnalisé. Une fois que vous avez ajouté un coût personnalisé et que vos appels comportent des données d’utilisation, vous pouvez récupérer les appels avec include_cost et consulter les coûts dans summary.weave.costs.
model = YourModel(attribute1="Hello", attribute2=1)
model.predict("world")

# Nous ajoutons ensuite un coût personnalisé à notre projet
weave_client.add_cost(
    llm_id="your_model_name", prompt_token_cost=0.1, completion_token_cost=0.2
)

# Nous pouvons ensuite interroger les appels, et avec include_costs=True
# nous recevons les coûts associés aux appels en retour
calls = weave_client.get_calls(filter={"trace_roots_only": True}, include_costs=True)

list(calls)
Vous disposez désormais d’un modèle personnalisé qui enregistre sa propre utilisation des jetons, d’un coût personnalisé enregistré pour le llm_id de ce modèle, ainsi que d’un moyen de récupérer des appels avec des données de coût par appel associées dans summary.weave.costs.