Add full tools support to the chat template

README.md

@@ -43,18 +43,32 @@ result = tokenizer.decode(out_tokens[0])
 print(result)
 ```
 
+## Preparing inputs with Hugging Face `transformers`
+
+```py
+from transformers import AutoTokenizer
+
+tokenizer = AutoTokenizer.from_pretrained("mistralai/Mixtral-8x22B-Instruct-v0.1")
+
+chat = [{"role": "user", "content": "Explain Machine Learning to me in a nutshell."}]
+
+tokens = tokenizer.apply_chat_template(chat, return_dict=True, return_tensors="pt", add_generation_prompt=True)
+```
+
 ## Inference with hugging face `transformers`
  
 ```py
 from transformers import AutoModelForCausalLM
- 
-model = AutoModelForCausalLM.from_pretrained("mistralai/Mixtral-8x22B-Instruct-v0.1")
+import torch
+
+# You can also use 8-bit or 4-bit quantization here
+model = AutoModelForCausalLM.from_pretrained("mistralai/Mixtral-8x22B-Instruct-v0.1", torch_dtype=torch.bfloat16, device_map="auto")
 model.to("cuda")
  
-generated_ids = model.generate(tokens, max_new_tokens=1000, do_sample=True)
+generated_ids = model.generate(**tokens, max_new_tokens=1000, do_sample=True)
 
-# decode with mistral tokenizer
-result = tokenizer.decode(generated_ids[0].tolist())
+# decode with HF tokenizer
+result = tokenizer.decode(generated_ids[0])
 print(result)
 ```
 
@@ -64,7 +78,7 @@ print(result)
 ---
 The Mixtral-8x22B-Instruct-v0.1 Large Language Model (LLM) is an instruct fine-tuned version of the [Mixtral-8x22B-v0.1](https://huggingface.co/mistralai/Mixtral-8x22B-v0.1).
 
-## Run the model 
+## Function calling example
 ```python
 from transformers import AutoModelForCausalLM
 from mistral_common.protocol.instruct.messages import (
@@ -122,56 +136,55 @@ sp_tokenizer = tokenizer_v3.instruct_tokenizer.tokenizer
 decoded = sp_tokenizer.decode(generated_ids[0])
 print(decoded)
 ```
-Alternatively, you can run this example with the Hugging Face tokenizer.
-To use this example, you'll need transformers version 4.39.0 or higher.
-```console
-pip install transformers==4.39.0
-```
+
+## Function calling with `transformers`
+
+To use this example, you'll need `transformers` version 4.42.0 or higher. Please see the 
+[function calling guide](https://huggingface.co/docs/transformers/main/chat_templating#advanced-tool-use--function-calling)
+in the `transformers` docs for more information.
+
 ```python
 from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
 
 model_id = "mistralai/Mixtral-8x22B-Instruct-v0.1"
 tokenizer = AutoTokenizer.from_pretrained(model_id)
-conversation=[
-    {"role": "user", "content": "What's the weather like in Paris?"},
-    {
-        "role": "tool_calls",
-        "content": [
-            {
-                "name": "get_current_weather",
-                "arguments": {"location": "Paris, France", "format": "celsius"},
-                
-            }
-        ]
-    },
-    {
-        "role": "tool_results",
-        "content": {"content": 22}
-    },
-    {"role": "assistant", "content": "The current temperature in Paris, France is 22 degrees Celsius."},
-    {"role": "user", "content": "What about San Francisco?"}
-]
-
-
-tools = [{"type": "function", "function": {"name":"get_current_weather", "description": "Get▁the▁current▁weather", "parameters": {"type": "object", "properties": {"location": {"type": "string", "description": "The city and state, e.g. San Francisco, CA"}, "format": {"type": "string", "enum": ["celsius", "fahrenheit"], "description": "The temperature unit to use. Infer this from the users location."}},"required":["location","format"]}}}]
+
+def get_current_weather(location: str, format: str):
+    """
+    Get the current weather
+
+    Args:
+        location: The city and state, e.g. San Francisco, CA
+        format: The temperature unit to use. Infer this from the users location. (choices: ["celsius", "fahrenheit"])
+    """
+    pass
+
+conversation = [{"role": "user", "content": "What's the weather like in Paris?"}]
+tools = [get_current_weather]
 
 # render the tool use prompt as a string:
 tool_use_prompt = tokenizer.apply_chat_template(
             conversation,
-            chat_template="tool_use",
             tools=tools,
             tokenize=False,
             add_generation_prompt=True,
-
 )
-model = AutoModelForCausalLM.from_pretrained("mistralai/Mixtral-8x22B-Instruct-v0.1")
 
 inputs = tokenizer(tool_use_prompt, return_tensors="pt")
 
-outputs = model.generate(**inputs, max_new_tokens=20)
+model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16, device_map="auto")
+
+outputs = model.generate(**inputs, max_new_tokens=1000)
 print(tokenizer.decode(outputs[0], skip_special_tokens=True))
 ```
 
+Note that, for reasons of space, this example does not show a complete cycle of calling a tool and adding the tool call and tool
+results to the chat history so that the model can use them in its next generation. For a full tool calling example, please
+see the [function calling guide](https://huggingface.co/docs/transformers/main/chat_templating#advanced-tool-use--function-calling), 
+and note that Mixtral **does** use tool call IDs, so these must be included in your tool calls and tool results. They should be
+exactly 9 alphanumeric characters.
+
 # Instruct tokenizer
 The HuggingFace tokenizer included in this release should match our own. To compare: 
 `pip install mistral-common`

tokenizer_config.json

@@ -6173,7 +6173,7 @@
     }
   },
   "bos_token": "<s>",
-  "chat_template": "{%- if messages[0]['role'] == 'system' %}\n    {%- set system_message = messages[0]['content'] %}\n    {%- set loop_messages = messages[1:] %}\n{%- else %}\n    {%- set loop_messages = messages %}\n{%- endif %}\n\n{{- bos_token }}\n{%- for message in loop_messages %}\n    {%- if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}\n        {{- raise_exception('After the optional system message, conversation roles must alternate user/assistant/user/assistant/...') }}\n    {%- endif %}\n    {%- if message['role'] == 'user' %}\n        {%- if loop.last and system_message is defined %}\n            {{- '[INST] ' + system_message + '\\n\\n' + message['content'] + '[/INST]' }}\n        {%- else %}\n            {{- '[INST] ' + message['content'] + '[/INST]' }}\n        {%- endif %}\n    {%- elif message['role'] == 'assistant' %}\n        {{- ' ' + message['content'] + eos_token}}\n    {%- else %}\n        {{- raise_exception('Only user and assistant roles are supported, with the exception of an initial optional system message!') }}\n    {%- endif %}\n{%- endfor %}\n",
+  "chat_template": "{%- if messages[0][\"role\"] == \"system\" %}\n    {%- set system_message = messages[0][\"content\"] %}\n    {%- set loop_messages = messages[1:] %}\n{%- else %}\n    {%- set loop_messages = messages %}\n{%- endif %}\n{%- if not tools is defined %}\n    {%- set tools = none %}\n{%- endif %}\n{%- set user_messages = loop_messages | selectattr(\"role\", \"equalto\", \"user\") | list %}\n\n{%- for message in loop_messages | rejectattr(\"role\", \"equalto\", \"tool\") | rejectattr(\"role\", \"equalto\", \"tool_results\") | selectattr(\"tool_calls\", \"undefined\") %}\n    {%- if (message[\"role\"] == \"user\") != (loop.index0 % 2 == 0) %}\n        {{- raise_exception(\"After the optional system message, conversation roles must alternate user/assistant/user/assistant/...\") }}\n    {%- endif %}\n{%- endfor %}\n\n{{- bos_token }}\n{%- for message in loop_messages %}\n    {%- if message[\"role\"] == \"user\" %}\n        {%- if tools is not none and (message == user_messages[-1]) %}\n            {{- \"[AVAILABLE_TOOLS] [\" }}\n            {%- for tool in tools %}\n        {%- set tool = tool.function %}\n        {{- '{\"type\": \"function\", \"function\": {' }}\n        {%- for key, val in tool.items() if key != \"return\" %}\n            {%- if val is string %}\n            {{- '\"' + key + '\": \"' + val + '\"' }}\n            {%- else %}\n            {{- '\"' + key + '\": ' + val|tojson }}\n            {%- endif %}\n            {%- if not loop.last %}\n            {{- \", \" }}\n            {%- endif %}\n        {%- endfor %}\n        {{- \"}}\" }}\n                {%- if not loop.last %}\n                    {{- \", \" }}\n                {%- else %}\n                    {{- \"]\" }}\n                {%- endif %}\n            {%- endfor %}\n            {{- \"[/AVAILABLE_TOOLS]\" }}\n            {%- endif %}\n        {%- if loop.last and system_message is defined %}\n            {{- \"[INST] \" + system_message + \"\\n\\n\" + message[\"content\"] + \"[/INST]\" }}\n        {%- else %}\n            {{- \"[INST] \" + message[\"content\"] + \"[/INST]\" }}\n        {%- endif %}\n    {%- elif message[\"role\"] == \"tool_calls\" or message.tool_calls is defined %}\n        {%- if message.tool_calls is defined %}\n            {%- set tool_calls = message.tool_calls %}\n        {%- else %}\n            {%- set tool_calls = message.content %}\n        {%- endif %}\n        {{- \"[TOOL_CALLS] [\" }}\n        {%- for tool_call in tool_calls %}\n            {%- set out = tool_call.function|tojson %}\n            {{- out[:-1] }}\n            {%- if not tool_call.id is defined or tool_call.id|length != 9 %}\n                {{- raise_exception(\"Tool call IDs should be alphanumeric strings with length 9!\") }}\n            {%- endif %}\n            {{- ', \"id\": \"' + tool_call.id + '\"}' }}\n            {%- if not loop.last %}\n                {{- \", \" }}\n            {%- else %}\n                {{- \"]\" + eos_token }}\n            {%- endif %}\n        {%- endfor %}\n    {%- elif message[\"role\"] == \"assistant\" %}\n        {{- \" \" + message[\"content\"] + eos_token}}\n    {%- elif message[\"role\"] == \"tool_results\" or message[\"role\"] == \"tool\" %}\n        {%- if message.content is defined and message.content.content is defined %}\n            {%- set content = message.content.content %}\n        {%- else %}\n            {%- set content = message.content %}\n        {%- endif %}\n        {{- '[TOOL_RESULTS] {\"content\": ' + content|string + \", \" }}\n        {%- if not message.tool_call_id is defined or message.tool_call_id|length != 9 %}\n            {{- raise_exception(\"Tool call IDs should be alphanumeric strings with length 9!\") }}\n        {%- endif %}\n        {{- '\"call_id\": \"' + message.tool_call_id + '\"}[/TOOL_RESULTS]' }}\n    {%- else %}\n        {{- raise_exception(\"Only user and assistant roles are supported, with the exception of an initial optional system message!\") }}\n    {%- endif %}\n{%- endfor %}\n",
   "clean_up_tokenization_spaces": false,
   "eos_token": "</s>",
   "legacy": false,