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

​
カスタムコストモデルの設定

Weave は、使用されたトークン数と使用された モデル に基づいてコストを計算します。 Weave は出力からこの使用量と モデル を取得し、それらをコールに関連付けます。 独自のトークン使用量を計算し、それを Weave に保存するシンプルなカスタム モデル を設定してみましょう。

​
環境 の設定

必要なパッケージをすべてインストールし、インポートします。 wandb.login() で簡単にログインできるように、環境 変数に WANDB_API_KEY を設定します（これは Colab の secret として提供される必要があります）。 これを ログ に記録したい W&B の プロジェクト を name_of_wandb_project に設定します。 注意: name_of_wandb_project は、トレース を ログ に記録する チーム を指定するために {team_name}/{project_name} の形式にすることもできます。 次に、 weave.init() を呼び出して Weave クライアントを取得します。 
%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)

​
Weave を使用した モデル の設定

from weave import Model

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

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

    # これは定義しているカスタム op です
    # 文字列を受け取り、使用量カウント、モデル名、および出力を含む辞書を出力します
    @weave.op()
    def custom_model_generate(self, input_data: str) -> dict:
        # ここにモデルのロジックを記述します
        # ここでカスタムの生成関数を使用します
        prediction = self.attribute1 + " " + input_data

        # 使用量カウント
        prompt_tokens = self.simple_token_count(input_data)
        completion_tokens = self.simple_token_count(prediction)

        # 使用量カウント、モデル名、および出力を保持する辞書を返します
        # Weave はこれを自動的にトレースに関連付けます
        # このオブジェクト {usage, model, output} は OpenAI コールの出力と一致します
        return {
            "usage": {
                "input_tokens": prompt_tokens,
                "output_tokens": completion_tokens,
                "total_tokens": prompt_tokens + completion_tokens,
            },
            "model": "your_model_name",
            "output": prediction,
        }

    # predict 関数内でカスタム生成関数を呼び出し、その出力を返します。
    @weave.op()
    def predict(self, input_data: str) -> dict:
        # ここでデータのポストプロセッシングを行うことができます
        outputs = self.custom_model_generate(input_data)
        return outputs["output"]

​
カスタムコストの追加

ここでカスタムコストを追加します。カスタムコストを設定し、コールに使用量が含まれているため、 include_cost を使用してコールを取得できるようになり、コールの summary.weave.costs の下にコストが表示されます。 
model = YourModel(attribute1="Hello", attribute2=1)
model.predict("world")

# 次に、プロジェクトにカスタムコストを追加します
weave_client.add_cost(
    llm_id="your_model_name", prompt_token_cost=0.1, completion_token_cost=0.2
)

# その後、コールをクエリできます。include_costs=True を指定すると
# コールに付随したコストを受け取ることができます
calls = weave_client.get_calls(filter={"trace_roots_only": True}, include_costs=True)

list(calls)