Uso de ferramentas

Modelos de chat são comumente treinados com suporte a “function-calling” ou “tool-use”. Ferramentas são funções fornecidas pelo usuário, que o modelo pode escolher chamar como parte de sua resposta. Por exemplo, os modelos poderiam ter acesso a uma ferramenta de calculadora para realizar operações aritméticas sem precisar executar o cálculo internamente.

Este guia demonstrará como definir ferramentas, como passá-las a um modelo de chat e como lidar com a saída do modelo quando ele chama uma ferramenta.

Passando ferramentas

Quando um modelo suporta o uso de ferramentas, passe as funções para o argumento tools de apply_chat_template(). As ferramentas são passadas como um JSON schema ou como funções Python. Se você passar funções Python, os argumentos, os tipos de argumentos e a docstring da função são analisados para gerar o JSON schema automaticamente.

Embora passar funções Python seja muito conveniente, o parser só consegue lidar com docstrings no estilo Google. Consulte os exemplos abaixo para saber como formatar uma função pronta para uso como ferramenta.

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]

Você pode adicionar opcionalmente um bloco Returns: à docstring e um tipo de retorno ao cabeçalho da função, mas a maioria dos modelos não usará essas informações. O parser também ignorará o código real dentro da função!

O que realmente importa é o nome da função, os nomes dos argumentos, os tipos dos argumentos e a docstring que descreve o propósito da função e o propósito de seus argumentos. Esses elementos criam a “assinatura” que o modelo usará para decidir se deve chamar a ferramenta.

Exemplo de chamada de ferramenta

Carregue um modelo e um tokenizador que suportem o uso de ferramentas, como NousResearch/Hermes-2-Pro-Llama-3-8B, mas você também pode considerar um modelo maior, como Command-R e Mixtral-8x22B, se seu hardware for capaz de suportá-lo.

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

Crie um histórico de chat.

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?"}
]

Em seguida, passe messages e uma lista de ferramentas para apply_chat_template(). Tokenize o chat e gere uma resposta.

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|>

O modelo de chat chamou a ferramenta get_current_temperature com os parâmetros corretos da docstring. Ele inferiu a França como a localização com base em Paris e que deveria usar Celsius para as unidades de temperatura.

Um modelo na verdade não consegue chamar a ferramenta por conta própria. Ele solicita uma chamada de ferramenta, e cabe a você processar essa chamada e adicioná-la, junto com o resultado, ao histórico do chat. Para modelos que suportam análise de resposta, a análise da resposta é processada automaticamente e você pode simplesmente usar parse_response() para extrair a chamada de ferramenta. Para outros modelos, você precisará traduzir manualmente a string de saída para um dicionário de chamada de ferramenta.

Independentemente da abordagem que você usar, a chamada de ferramenta deve estar na chave tool_calls de uma mensagem assistant. Esta é a API recomendada e deve ser suportada pelo modelo de chat da maioria dos modelos que utilizam ferramentas.

Embora tool_calls seja semelhante à API da OpenAI, a API da OpenAI usa uma string JSON como formato de tool_calls. Isso pode causar erros ou comportamento estranho do modelo se for usado no Transformers, que espera um dict.

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

Anexe a resposta da ferramenta ao histórico do chat com o papel tool.

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

Por fim, permita que o modelo leia a resposta da ferramenta e responda ao usuário.

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|>

Embora a chave na mensagem do assistente seja chamada de tool_calls, na maioria dos casos, os modelos emitem apenas uma chamada de ferramenta por vez. Alguns modelos mais antigos emitem múltiplas chamadas de ferramenta ao mesmo tempo, mas esse é um processo significativamente mais complexo, pois é necessário lidar com múltiplas respostas de ferramentas simultaneamente e desambiguá-las, frequentemente usando IDs de chamada de ferramenta. Consulte o cartão do modelo para ver exatamente qual formato o modelo espera para as chamadas de ferramenta.

Esquemas JSON

Outra forma de definir ferramentas é passando um JSON schema.

Você também pode chamar manualmente as funções de baixo nível que convertem funções Python em esquemas JSON e, em seguida, verificar ou editar os esquemas gerados. Isso geralmente não é necessário, mas é útil para entender a mecânica subjacente. É particularmente importante para os autores de modelos de chat que precisam acessar o esquema JSON para renderizar as definições de ferramentas.

O método apply_chat_template() usa a função get_json_schema para converter chamáveis do Python em um esquema JSON. Isso inclui métodos: self e cls são tratados como argumentos de receptor implícitos e são ignorados.

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"]
    }
  }
}

Não entraremos nos detalhes do próprio esquema JSON aqui, pois ele já está muito bem documentado em outros lugares. No entanto, mencionaremos que você pode passar dicionários de esquema JSON para o argumento tools de apply_chat_template() em vez de funções Python:

# 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]
)