ツール使用

Chat modelsは、“function-calling”または“tool-use”のサポート付きで学習されることが一般的です。ツールとは、ユーザーが提供する関数のことで、モデルは応答の一部としてこれらを呼び出すかどうかを選択できます。例えば、モデルは計算機ツールにアクセスすることで、内部で計算を実行せずに算術演算を行うことができます。

このガイドでは、ツールの定義方法、チャットモデルへのツールの渡し方、およびモデルがツールを呼び出した場合の出力の処理方法について説明します。

ツールのパス

モデルがツール使用をサポートしている場合、apply_chat_template() の tools 引数に関数を渡します。ツールは、JSON スキーマまたは Python 関数として渡します。Python 関数を渡す場合、引数、引数の型、および関数の docstring が解析され、JSON スキーマが自動的に生成されます。

Python関数を渡すことは非常に便利ですが、パーサーはGoogleスタイルのdocstringのみを処理できます。ツール対応の関数の書式付け方法については、以下の例を参照してください。

def get_current_temperature(location: str, unit: str):
    """
    Get the current temperature at a location.

    Args:
        location: The location to get the temperature for, in the format "City, Country"
        unit: The unit to return the temperature in. (choices: ["celsius", "fahrenheit"])
    """
    return 22.  # A real function should probably actually get the temperature!

def get_current_wind_speed(location: str):
    """
    Get the current wind speed in km/h at a given location.

    Args:
        location: The location to get the wind speed for, in the format "City, Country"
    """
    return 6.  # A real function should probably actually get the wind speed!

tools = [get_current_temperature, get_current_wind_speed]

ドキュメント文字列にReturns:ブロックを追加し、関数ヘッダーに戻り値の型を指定することもできますが、ほとんどのモデルはこの情報を使用しません。パーサーは関数内の実際のコードも無視します。

本当に重要なのは、関数名、引数名、引数の型、および関数の目的と引数の目的を説明するドキュメント文字列です。これらが、モデルがツールを呼び出すかどうかを判断するために使用する「シグネチャ」を構成します。

ツール呼び出しの例

ツール使用に対応したモデルとトークナイザー、例えば NousResearch/Hermes-2-Pro-Llama-3-8B を読み込みます。ハードウェアが対応している場合は、Command-R や Mixtral-8x22B のようなより大きなモデルも検討できます。

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

checkpoint = "NousResearch/Hermes-2-Pro-Llama-3-8B"
tokenizer = AutoTokenizer.from_pretrained(checkpoint)
model = AutoModelForCausalLM.from_pretrained(checkpoint, dtype="auto", device_map="auto")

チャット履歴を作成します。

messages = [
  {"role": "system", "content": "You are a bot that responds to weather queries. You should reply with the unit used in the queried location."},
  {"role": "user", "content": "Hey, what's the temperature in Paris right now?"}
]

次に、messages とツールのリストを apply_chat_template() に渡します。チャットをトークン化し、応答を生成します。

inputs = tokenizer.apply_chat_template(messages, tools=tools, add_generation_prompt=True, return_dict=True, return_tensors="pt")
outputs = model.generate(**inputs.to(model.device), max_new_tokens=128)
print(tokenizer.decode(outputs[0][len(inputs["input_ids"][0]):]))

<tool_call>
{"arguments": {"location": "Paris, France", "unit": "celsius"}, "name": "get_current_temperature"}
</tool_call><|im_end|>

チャットモデルは、get_current_temperature ツールを、docstring に記載された正しいパラメータで呼び出しました。パリに基づいて場所をフランスと推論し、温度の単位としてセルシウスを使うべきだと判断しました。

モデルはツールを直接呼び出すことはできません。ツール呼び出しを要求し、その呼び出しと結果をチャット履歴に追加する処理は、あなたが行う必要があります。レスポンスの解析をサポートするモデルでは、解析が自動的に処理され、parse_response()を使用してツール呼び出しを抽出できます。その他のモデルでは、出力文字列を手動でツール呼び出しの辞書に変換する必要があります。

使用するアプローチにかかわらず、ツール呼び出しはassistantメッセージのtool_callsキーに配置する必要があります。これが推奨されるAPIであり、ツールを使用するほとんどのモデルのチャットテンプレートでサポートされるべきです。

tool_calls は OpenAI API と似ていますが、OpenAI API は tool_calls の形式として JSON 文字列を使用します。dict を期待する Transformers でこれを使用すると、エラーやモデルの奇妙な動作を引き起こす可能性があります。

tool_call = {"name": "get_current_temperature", "arguments": {"location": "Paris, France", "unit": "celsius"}}
messages.append({"role": "assistant", "tool_calls": [{"type": "function", "function": tool_call}]})

toolロールでツールレスポンスをチャット履歴に追加します。

messages.append({"role": "tool", "content": "22"})  # Note that the returned content is always a string!

最後に、モデルがツールの応答を読み取り、ユーザーに返信できるようにします。

inputs = tokenizer.apply_chat_template(messages, tools=tools, add_generation_prompt=True, return_dict=True, return_tensors="pt")
out = model.generate(**inputs.to(model.device), max_new_tokens=128)
print(tokenizer.decode(out[0][len(inputs["input_ids"][0]):]))

The temperature in Paris, France right now is 22°C.<|im_end|>

アシスタントメッセージのキーはtool_callsと呼ばれますが、ほとんどの場合、モデルは一度に1つのツール呼び出しのみを出力します。一部の古いモデルは同時に複数のツール呼び出しを出力しますが、これは複数のツール応答を同時に処理し、ツール呼び出しIDを使用してそれらを区別する必要があるため、はるかに複雑なプロセスです。モデルがツール呼び出しに対してどのような形式を期待しているかについては、モデルカードを参照してください。

JSONスキーマ

ツールを定義するもう一つの方法は、JSON スキーマを渡すことです。

Python関数をJSONスキーマに変換する低レベル関数を手動で呼び出し、生成されたスキーマを確認または編集することもできます。通常これは必要ありませんが、内部の仕組みを理解するには有用です。ツール定義をレンダリングするためにJSONスキーマへのアクセスが必要なチャットテンプレートの作成者にとって、特に重要です。

apply_chat_template() メソッドは、get_json_schema 関数を使用して、Python の呼び出し可能オブジェクトを JSON スキーマに変換します。これにはメソッドが含まれます。self および cls は暗黙のレシーバー引数として扱われ、無視されます。

from transformers.utils import get_json_schema

def multiply(a: float, b: float):
    """
    A function that multiplies two numbers

    Args:
        a: The first number to multiply
        b: The second number to multiply
    """
    return a * b

schema = get_json_schema(multiply)
print(schema)

{
  "type": "function",
  "function": {
    "name": "multiply",
    "description": "A function that multiplies two numbers",
    "parameters": {
      "type": "object",
      "properties": {
        "a": {
          "type": "number",
          "description": "The first number to multiply"
        },
        "b": {
          "type": "number",
          "description": "The second number to multiply"
        }
      },
      "required": ["a", "b"]
    }
  }
}

JSONスキーマ自体の詳細については、すでに他の場所で非常に詳細にドキュメント化されているため、ここでは省略します。ただし、Python関数の代わりに、apply_chat_template()のtools引数にJSONスキーマの辞書を渡すことができる点には言及しておきます。

# A simple function that takes no arguments
current_time = {
  "type": "function",
  "function": {
    "name": "current_time",
    "description": "Get the current local time as a string.",
    "parameters": {
      'type': 'object',
      'properties': {}
    }
  }
}

# A more complete function that takes two numerical arguments
multiply = {
  'type': 'function',
  'function': {
    'name': 'multiply',
    'description': 'A function that multiplies two numbers',
    'parameters': {
      'type': 'object',
      'properties': {
        'a': {
          'type': 'number',
          'description': 'The first number to multiply'
        },
        'b': {
          'type': 'number', 'description': 'The second number to multiply'
        }
      },
      'required': ['a', 'b']
    }
  }
}

model_input = tokenizer.apply_chat_template(
    messages,
    tools = [current_time, multiply]
)

Update on GitHub