W&B Weave を使ってLLMを評価してみる

W&B Weave とは

https://wandb.ai/site/weave/
https://weave-docs.wandb.ai/

Weights & Biases が開発した LLM の評価や分析を行うツールです。LLM を使ったアプリケーションを開発するときや、研究でLLMの評価や分析を行うときにも役立つツールのように思えます。2024年の4月にアナウンスされた かなり新しいプロダクトですね。

詳しいことは公式のドキュメントを読むのが良いとは思いますが、ドキュメントは OpenAI API を使ったものなので、今回はローカル環境で動く Gemma2 を使って Weave を試してみます。

インストール

まずはパッケージをインストールします。

pip install weave

W&B のアカウントを持っていない場合は作成します。その後、W&B の公式ドキュメント に従って、プログラムを実行するマシンでログインします。 weave をグローバルにインストールしていない場合は pipx か何かでグローバルにインストールしましょう。（weave をインストールすると wandb もインストールされますが、グローバル環境で必要なのは wandb だけなので、実際にはそれだけインストールすれば十分です。 ）
あるいは API key を設定します

実装

今回は単語を音節に分解してみます。簡単にデータセットを作成して、（このタスクやデータセットの妥当性はさておき、）Gemma2 を評価していきます。
評価用のプログラムの完成形は以下のようになります。

評価用プログラム

import asyncio

import torch
import weave
from transformers import (
    Gemma2ForCausalLM,
    GemmaTokenizerFast,
    PreTrainedModel,
    PreTrainedTokenizerFast,
)


class SeparateSyllablesModel(weave.Model):
    prompt_template: str
    chat: list[dict]
    model_id: str
    _tokenizer: PreTrainedTokenizerFast
    _model: PreTrainedModel

    def __init__(self, /, **data):
        super().__init__(**data)
        self._model = Gemma2ForCausalLM.from_pretrained(
            self.model_id,
            torch_dtype=torch.bfloat16,
            device_map="auto",
        )
        self._tokenizer = GemmaTokenizerFast.from_pretrained(self.model_id)

    @weave.op
    def predict(self, input_text):
        query = self.prompt_template.format(input_text)

        current_chat = self.chat.copy()
        current_chat.append({"role": "user", "content": query})

        prompt = self._tokenizer.apply_chat_template(
            current_chat, tokenize=False, add_generation_prompt=True
        )
        input_ids = self._tokenizer.encode(
            prompt, add_special_tokens=False, return_tensors="pt"
        ).to(self._model.device)

        outputs = self._model.generate(input_ids, max_new_tokens=1024)
        decoded_outputs = self._tokenizer.decode(outputs[0])

        response = decoded_outputs.split("<start_of_turn>model")[-1].strip()
        response = response.split("<end_of_turn>")[0].strip()
        response = response.split("\n")[0]  # 1行目だけを見る

        return response


def build_dataset():
    return [
        {"id": 1, "input_text": "weight", "syllables": "weight"},
        {"id": 2, "input_text": "and", "syllables": "and"},
        {"id": 3, "input_text": "bias", "syllables": "bi-as"},
        {"id": 4, "input_text": "evaluation", "syllables": "e-val-u-a-tion"},
    ]


@weave.op()
def counted_mora_score(syllables: str, model_output: str):
    return {"correct": syllables == model_output}


def main():
    prompt_path = "src/prompt_text.txt"
    prompt_template = open(prompt_path, "r").read()
    model_id = "google/gemma-2-9b-it"

    weave.init("count-mora")
    chat = [
        {"role": "user", "content": prompt_template.format("hello")},
        {"role": "model", "content": "he-llo"},
    ]
    model = SeparateSyllablesModel(
        prompt_template=prompt_template,
        model_id=model_id,
        chat=chat,
    )
    evaluation = weave.Evaluation(
        dataset=build_dataset(),
        scorers=[counted_mora_score],
    )
    print(asyncio.run(evaluation.evaluate(model)))


if __name__ == "__main__":
    main()

主な構成要素は

1. SeparateSyllablesModel
2. counted_mora_score

の2つです。つまり推論するモデルとその推論結果の評価関数です。 SeparateSyllablesModel に渡している3つのパラメータによって、 SeparateSyllablesModel がバージョニングされるので、パラメータを変えたときの性能の比較が簡単になります。

Model

あらためて実装を見てみます

class SeparateSyllablesModel(weave.Model):
    prompt_template: str
    chat: list[dict]
    model_id: str
    _tokenizer: PreTrainedTokenizerFast
    _model: PreTrainedModel

    def __init__(self, /, **data):
        super().__init__(**data)
        self._model = Gemma2ForCausalLM.from_pretrained(
            self.model_id,
            torch_dtype=torch.bfloat16,
            device_map="auto",
        )
        self._tokenizer = GemmaTokenizerFast.from_pretrained(self.model_id)

    @weave.op
    def predict(self, input_text):
        query = self.prompt_template.format(input_text)

        current_chat = self.chat.copy()
        current_chat.append({"role": "user", "content": query})

        prompt = self._tokenizer.apply_chat_template(
            current_chat, tokenize=False, add_generation_prompt=True
        )
        input_ids = self._tokenizer.encode(
            prompt, add_special_tokens=False, return_tensors="pt"
        ).to(self._model.device)

        outputs = self._model.generate(input_ids, max_new_tokens=1024)
        decoded_outputs = self._tokenizer.decode(outputs[0])

        response = decoded_outputs.split("<start_of_turn>model")[-1].strip()
        response = response.split("<end_of_turn>")[0].strip()
        response = response.split("\n")[0]  # 1行目だけを見る

        return response

weave.Evaluation を使う場合、 predict() の引数になっている input_text には、weave.Evaluation に渡した dataset の中の同名のフィールドの値が渡されます。
今回の場合だと、dataset 中の input_text の値（weight や bias）が predict() に渡されます。

def build_dataset():
    return [
        {"id": 1, "input_text": "weight", "syllables": "weight"},
        {"id": 2, "input_text": "and", "syllables": "and"},
        {"id": 3, "input_text": "bias", "syllables": "bi-as"},
        {"id": 4, "input_text": "evaluation", "syllables": "e-val-u-a-tion"},
    ]

evaluation = weave.Evaluation(
        dataset=build_dataset(),
        scorers=[counted_mora_score],
    )

predict() の中身については Gemma2 の話なので割愛します。

ちなみに、ここで作成した Model は Weave 上で以下のように見ることができます。

Scorer

@weave.op()
def counted_mora_score(syllables: str, model_output: str):
    return {"correct": syllables == model_output}

評価用の関数です。引数の syllables には dataset の中の syllables フィールドの値が渡されます。 model_output は名前の通り、Model （の predict）の出力が渡されます。

ここでは完全一致によって評価していますが、もちろんその他の評価指標を用いて評価することもできます。
精度や再現率、F値を計算してくれる MultiTaskBinaryClassificationF1 もデフォルトで提供されているようなので、痒いところに手が届く感じもします。

Weave で結果を見る・比較する

結果を見る

Weave 上では以下の用に結果を確認することができます。
先ほど実装した SeparateSyllablesModel (SeparateSyllablesModel:v0）は one-shot でしたが、 two-shot にした SeparateSyllablesModel:v1 も作成しました。細かいところは置いておいて、この2つの設定の Evaluation の結果を見てみます。

先ほど、モデルの出力を正解との完全一致によって評価するようにしたため、完全一致だったデータの数（true_count）とそのデータセット全体に占める割合（true_fraction）がひと目で分かるようになっています。

それぞれの Evaluation の詳細画面は以下の通りです。各データの入力と正解、実際の出力が表示されています。もちろん、結果は CSV や JSON などでエクスポートできます。API経由で直接取得することもできます。

比較する

one-shot と two-shot の Evaluation の結果が大体わかったところで、この2つを比較してみます。Weave では以下のようにいくつかの Evaluation を比較して、可視化する機能が提供されています。

Weave では全体の精度の比較だけではなく、それぞれの入力に対する出力を比較することができます。

まとめ

今回は基本的な機能に絞って、Weave を使って LLM を評価してみました。Weave を使うことで、お手軽にLLMの評価や分析をすることができます。副産物として、 Weave が推論結果や評価結果、それらの実行時のパラメータをまとめてくれているので、実験データの管理を丸投げすることもできます。（もちろん大切な実験データはバックアップを取っておくべきですが……）

個人的には Huggingface でモデルやデータセットを公開しているように、論文の実験結果なんかも Weave のようなプラットフォームで公開してくれたほうが、より開かれた研究になるのでは、と思ったりもしました。

今回は紹介しませんでしたが、OpenAI API を使う場合には token 数の計算やコスト計算までやってくれるので、それらの見積もりや記録が容易にできます。これは個人的に地味に嬉しい機能です。

その他にもまだまだ面白そうな機能や便利機能があるので、みなさんも一度試してみてはいかがでしょうか。