Entendiendo los Tokenizers de Mistral

[!IMPORTANT] Este documento está obsoleto. Por favor, consulta la documentación completa y detallada aquí.

Desde el tokenizer original V1 hasta los más recientes tokenizers V3 y Tekken, los tokenizers de Mistral han experimentado cambios sutiles relacionados con la forma de tokenizar para los modelos de instrucción. Hemos avanzado mucho en optimización e investigación para encontrar el mejor enfoque. Hoy, profundizaremos en estos tokenizers, desmitificaremos cualquier fuente de debate y exploraremos cómo funcionan, las plantillas de chat adecuadas para usar con cada uno y su historia dentro de la comunidad.

En este artículo, nos centraremos principalmente en la tokenización de instrucciones y las plantillas de chat para el seguimiento de instrucciones simples. No profundizaremos en la llamada a funciones ni en el "fill in the middle" (FIM).

La información fundamental para todo lo que se presenta aquí se puede encontrar explorando los repositorios en Hugging Face, así como en GitHub, específicamente mistral_common.

"<s> [INST] user message [/INST] assistant message</s> [INST] new user message [/INST]"

"<s>[INST] user message[/INST] assistant message</s>[INST] new user message[/INST]"

"<s>[SUFFIX]suffix[PREFIX] prefix"

"<s>[INST] user message[/INST] assistant message</s>[INST] new user message[/INST]"

"<s>[SUFFIX]suffix[PREFIX] prefix"

"<s>[INST]user message[/INST]assistant message</s>[INST]new user message[/INST]"

Tokenizer V1

Las primeras versiones, Mistral 7B V1 y V2, así como Mixtral 8x7B V1, fueron recibidas con mucho aprecio y cariño por parte de la comunidad. Sin embargo, esto no impidió algunos desacuerdos y debates dentro de la comunidad con respecto a las plantillas de chat correctas para los tokenizers de instrucción. Hoy, puedes confiar de manera fiable en mistral_common como la verdad fundamental para el proceso de tokenización, pero en aquel entonces, la mayor parte de la comunidad se basaba en las plantillas de chat Jinja disponibles.

La comunidad fue muy rápida en notar inconsistencias en torno a los tokenizers y las plantillas de chat, logrando mejores resultados con versiones personalizadas que ajustaban los espacios en blanco alrededor de las cadenas especiales. ¡Es increíble lo rápida y útil que es la comunidad para detectar cualquier problema y error relacionado con los modelos!

El tokenizer utilizado para estas primeras versiones, basado en sentencepiece, utiliza una plantilla similar a la de los antiguos modelos Llama. Este tokenizer maneja los mensajes de instrucción de la siguiente manera:

Plantilla de Chat de Instrucción

La plantilla de chat era, como se mencionó anteriormente, muy similar a la primera plantilla de Llama, y resultaría en cadenas como la siguiente.

<s> [INST] user message [/INST] assistant message</s> [INST] new user message [/INST]

Aquí, las únicas cadenas especiales eran [INST] para iniciar el mensaje del usuario y [/INST] para finalizar el mensaje del usuario, dando paso a la respuesta del asistente. El BOS (inicio de cadena) estaba y sigue estando representado con <s>, y el EOS (fin de cadena) es </s>, utilizado al final de cada finalización, terminando cualquier mensaje del asistente.

Los espacios en blanco son de extrema importancia.

Ten en cuenta que <s> y </s> son más como cadenas que representan el BOS y el EOS que cadenas reales. Sus IDs correspondientes son 1 y 2.

⚠️ Esta plantilla siempre hará que el modelo genere un token que comienza con un espacio en blanco, siempre teniendo un espacio inicial. Dependiendo de tu motor de inferencia, puede que elimine o no ese espacio en blanco por defecto. Por lo tanto, es esencial asegurarse de que un espacio en blanco (y solo uno) esté presente después de cada respuesta del asistente al usar la plantilla.

La plantilla Jinja para el transformers de HuggingFace para este tipo de cadena podría verse así:

{{ bos_token }}{% for message in messages %}{% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}{{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}{% endif %}{% if message['role'] == 'user' %}{{ ' [INST] ' + message['content'] + ' [/INST]' }}{% elif message['role'] == 'assistant' %}{{ ' ' + message['content'] + eos_token}}{% else %}{{ raise_exception('Only user and assistant roles are supported!') }}{% endif %}{% endfor %}

Puedes ver la que se está usando actualmente por defecto echando un vistazo al tokenizer_config.json.

Aquí está formateada para que sea más fácil de leer, lo que facilita la comprensión de cómo funcionan:

{{ bos_token }}
{% for message in messages %}
    {% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}
        {{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}
    {% endif %}
    {% if message['role'] == 'user' %}
        {{ ' [INST] ' + message['content'] + ' [/INST]' }}
    {% elif message['role'] == 'assistant' %}
        {{ ' ' + message['content'] + eos_token }}
    {% else %}
        {{ raise_exception('Only user and assistant roles are supported!') }}
    {% endif %}
{% endfor %}

⚠️ Como puedes ver, cada message['content'] siempre va precedido de un espacio en blanco. Esto asume que message['content'] tiene el primer espacio en blanco eliminado manualmente por ti o por tu biblioteca de inferencia.

Si quieres aplicar la plantilla manualmente, puedes hacerlo en Python:

from jinja2 import Template, Environment, exceptions

JINJA_TEMPLATE = "{{ bos_token }}{% for message in messages %}{% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}{{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}{% endif %}{% if message['role'] == 'user' %}{{ ' [INST] ' + message['content'] + ' [/INST]' }}{% elif message['role'] == 'assistant' %}{{ ' ' + message['content'] + eos_token}}{% else %}{{ raise_exception('Only user and assistant roles are supported!') }}{% endif %}{% endfor %}"

def raise_exception(message):
    raise ValueError(message)

def apply_jinja_template(messages, bos_token='<s>', eos_token='</s>'):
    template_str = JINJA_TEMPLATE

    env = Environment()
    env.globals['raise_exception'] = raise_exception

    template = env.from_string(template_str)

    try:
        result = template.render(messages=messages, bos_token=bos_token, eos_token=eos_token)
        return result
    except exceptions.TemplateError as e:
        raise ValueError(f"Template rendering error: {e}")

print(apply_jinja_template(messages))

Lógica de Tokenización de Instrucciones

Ahora profundicemos en la lógica original detrás del proceso de tokenización que puedes ver en mistral_common. A menudo no simplemente proporcionarías el texto al tokenizer y codificarías toda la cadena tal como está, mistral-common en realidad va directamente de request -> int, mientras que otros enfoques como el utilizado por transformers están más cerca de request -> str -> int. La forma en que se hace realmente para el proceso de tokenización es más cercana a la siguiente lógica:

• Tokenizar cada mensaje:
  Tendrías mensajes de usuario y mensajes de asistente. Los tomarías por separado como cadenas individuales, con el mensaje del usuario encapsulado con las cadenas especiales, y los codificarías. En el ejemplo anterior, esto sería algo como:
  encode("[INST] user message [/INST]"), encode("assistant message"), y encode("[INST] new user message [/INST]"), y así sucesivamente...

• Concatenar:
  Una vez que cada mensaje es tokenizado y codificado, los concatenarías todos, anteponiéndoles el ID de BOS y separándolos por pares con el ID de EOS:
  BOS_ID + encode("[INST] user message [/INST]") + encode("assistant message") + EOS_ID + encode("[INST] new user message [/INST]")

Podrías notar que si se sigue paso a paso, esto en realidad coincidiría con una cadena y espacios en blanco añadidos de la siguiente manera:

<s>[INST]_user message_[/INST]_assistant message</s>[INST]_new user message_[/INST]

_{^{"_" representa los espacios en blanco añadidos al encapsular con las cadenas especiales}}

Si se compara con la plantilla de chat mencionada anteriormente, esto omite algunos espacios en blanco después del BOS y el EOS. ¿De dónde vienen entonces? Vienen de sentencepiece al codificar. Por defecto, comenzará con un espacio en blanco inicial. Así que cada vez que encode, estamos añadiendo un espacio en blanco inicial. Por lo tanto:

• encode("[INST] user message [/INST]") en realidad está codificando encode(" [INST] user message [/INST]")

Esto añade los espacios en blanco iniciales, haciendo que la cadena final se vea más como:

<s>_[INST] user message [/INST] assistant message</s>_[INST] new user message [/INST]

_{^{"_" representa los nuevos espacios en blanco añadidos por sentencepiece}}

Finalmente, token por token dará lo siguiente: