Uso de herramientas

Los modelos de chat suelen entrenarse con soporte para “function-calling” o “tool-use”. Las herramientas son funciones proporcionadas por el usuario, que el modelo puede elegir llamar como parte de su respuesta. Por ejemplo, los modelos podrían tener acceso a una herramienta de calculadora para realizar operaciones aritméticas sin tener que realizar el cálculo internamente.

Esta guía mostrará cómo definir herramientas, cómo pasarlas a un modelo de chat y cómo manejar la salida del modelo cuando llama a una herramienta.

Pasar herramientas

Cuando un modelo admite el uso de herramientas, pase funciones al argumento tools de apply_chat_template(). Las herramientas se pasan como un JSON schema o como funciones de Python. Si pasa funciones de Python, los argumentos, los tipos de argumentos y la docstring de la función se analizan para generar automáticamente el JSON schema.

Aunque pasar funciones de Python es muy conveniente, el analizador solo puede manejar docstrings en estilo Google. Consulte los ejemplos a continuación para ver cómo formatear una función lista para usar como herramienta.

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]

Puedes añadir opcionalmente un bloque Returns: a la docstring y un tipo de retorno a la cabecera de la función, pero la mayoría de los modelos no utilizarán esta información. El parser también ignorará el código real dentro de la función.

Lo que realmente importa es el nombre de la función, los nombres de los argumentos, los tipos de argumentos y la docstring que describe el propósito de la función y el propósito de sus argumentos. Estos elementos crean la “firma” que el modelo utilizará para decidir si llamar a la herramienta.

Ejemplo de invocación de herramientas

Carga un modelo y un tokenizador que admitan el uso de herramientas, como NousResearch/Hermes-2-Pro-Llama-3-8B, aunque también puedes considerar un modelo más grande como Command-R y Mixtral-8x22B si tu hardware lo permite.

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

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

A continuación, pase messages y una lista de herramientas a apply_chat_template(). Tokenice el chat y genere una respuesta.

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

El modelo de chat llamó a la herramienta get_current_temperature con los parámetros correctos de la docstring. Inferió Francia como la ubicación en base a París, y que debía usar Celsius para las unidades de temperatura.

Un modelo no puede llamar a la herramienta por sí mismo. Solicita una llamada a la herramienta, y es tu responsabilidad gestionar esa llamada y añadir tanto la llamada como el resultado al historial del chat. Para los modelos que admiten response parsing, el análisis de la respuesta se gestiona automáticamente y solo tienes que usar parse_response() para extraer la llamada a la herramienta. Para otros modelos, tendrás que traducir manualmente la cadena de salida a un diccionario de llamada a la herramienta.

Independientemente del enfoque que utilice, la llamada a la herramienta debe ir en la clave tool_calls de un mensaje de assistant. Esta es la API recomendada y debería ser compatible con la plantilla de chat de la mayoría de los modelos que utilizan herramientas.

Aunque tool_calls es similar a la API de OpenAI, la API de OpenAI utiliza una cadena JSON como formato de tool_calls. Esto puede provocar errores o un comportamiento extraño del modelo si se usa en Transformers, que espera un diccionario.

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

Agrega la respuesta de la herramienta al historial del chat con el rol tool.

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

Por último, permite que el modelo lea la respuesta de la herramienta y responda al usuario.

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

Aunque la clave en el mensaje del asistente se llama tool_calls, en la mayoría de los casos, los modelos solo emiten una llamada a herramienta a la vez. Algunos modelos más antiguos emiten múltiples llamadas a herramienta al mismo tiempo, pero este es un proceso significativamente más complejo, ya que necesitas manejar múltiples respuestas de herramientas a la vez y desambiguarlas, a menudo usando IDs de llamada a herramienta. Consulta la tarjeta del modelo para ver exactamente qué formato espera el modelo para las llamadas a herramienta.

Esquemas JSON

Otra forma de definir herramientas es pasando un esquema JSON.

También puede llamar manualmente a las funciones de bajo nivel que convierten funciones de Python en esquemas JSON y, a continuación, verificar o editar los esquemas generados. Esto generalmente no es necesario, pero resulta útil para comprender la mecánica subyacente. Es especialmente importante para los autores de plantillas de chat que necesitan acceder al esquema JSON para renderizar las definiciones de herramientas.

El método apply_chat_template() utiliza la función get_json_schema para convertir los llamables de Python en un esquema JSON. Esto incluye los métodos: self y cls se tratan como argumentos de receptor implícitos y se ignoran.

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

No entraremos en los detalles del propio esquema JSON aquí, ya que está muy bien documentado en otros lugares. Sin embargo, mencionaremos que puede pasar diccionarios de esquemas JSON al argumento tools de apply_chat_template() en lugar de funciones de 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]
)