Add processor chat template

README.md

@@ -23,9 +23,10 @@ PROMPT = "<s>[INST]Describe the images.\n[IMG][IMG][IMG][IMG][/INST]"
 
 inputs = processor(text=PROMPT, images=IMG_URLS, return_tensors="pt").to("cuda")
 generate_ids = model.generate(**inputs, max_new_tokens=500)
-ouptut = processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+output = processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
 ```
-I got something like this:
+
+You should get an output similar to the below:
 ```
 
 """
@@ -50,4 +51,61 @@ Sure, let's break down each image description:
 
 Each image captures a different scene, from a close-up of a dog to expansive natural landscapes, showcasing various elements of nature and human interaction with it.
 """
-```
+```
+
+You can also use a chat template to format your chat history for Pixtral. Make sure that the `images` argument to the `processor` contains the images in the order
+that they appear in the chat, so that the model understands where each image is supposed to go.
+
+Here's an example with text and multiple images interleaved in the same message:
+
+```python
+from PIL import Image
+from transformers import AutoProcessor, LlavaForConditionalGeneration
+model_id = "mistral-community/pixtral-12b"
+model = LlavaForConditionalGeneration.from_pretrained(model_id)
+processor = AutoProcessor.from_pretrained(model_id)
+
+url_dog = "https://picsum.photos/id/237/200/300"
+url_mountain = "https://picsum.photos/seed/picsum/200/300"
+
+chat = [
+    {
+      "role": "user", "content": [
+        {"type": "text", "content": "Can this animal"}, 
+        {"type": "image"}, 
+        {"type": "text", "content": "live here?"}, 
+        {"type": "image"}
+      ]
+    }
+]
+
+prompt = processor.apply_chat_template(chat)
+inputs = processor(text=prompt, images=[url_dog, url_mountain], return_tensors="pt").to(model.device)
+generate_ids = model.generate(**inputs, max_new_tokens=500)
+output = processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+```
+
+You should get something like this:
+
+```
+Can this animallive here?Certainly! Here are some details about the images you provided:
+
+### First Image
+- **Description**: The image shows a black dog lying on a wooden surface. The dog has a curious expression with its head tilted slightly to one side.
+- **Details**: The dog appears to be a young puppy with soft, shiny fur. Its eyes are wide and alert, and it has a playful demeanor.
+- **Context**: This image could be used to illustrate a pet-friendly environment or to showcase the dog's personality.
+
+### Second Image
+- **Description**: The image depicts a serene landscape with a snow-covered hill in the foreground. The sky is painted with soft hues of pink, orange, and purple, indicating a sunrise or sunset.
+- **Details**: The hill is covered in a blanket of pristine white snow, and the horizon meets the sky in a gentle curve. The scene is calm and peaceful.
+- **Context**: This image could be used to represent tranquility, natural beauty, or a winter wonderland.
+
+### Combined Context
+If you're asking whether the dog can "live here," referring to the snowy landscape, it would depend on the breed and its tolerance to cold weather. Some breeds, like Huskies or Saint Bernards, are well-adapted to cold environments, while others might struggle. The dog in the first image appears to be a breed that might prefer warmer climates.
+
+Would you like more information on any specific aspect?
+```
+
+While it may appear that spacing in the input is disrupted, this is caused by us skipping special tokens for display, and actually "Can this animal" and "live here" are
+correctly separated by image tokens. Try decoding with special tokens included to see exactly what the model sees!
+

chat_template.json

@@ -0,0 +1,3 @@
+{
+  "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\" }}\n        {%- else %}\n            {{- \"[INST]\" }}\n        {%- endif %}\n        {%- if message[\"content\"] is not string %}\n            {%- for chunk in message[\"content\"] %}\n                {%- if chunk[\"type\"] == \"text\" %}\n                    {{- chunk[\"content\"] }}\n                {%- elif chunk[\"type\"] == \"image\" %}\n                    {{- \"[IMG]\" }}\n                {%- else %}\n                    {{- raise_exception(\"Unrecognized content type!\") }}\n                {%- endif %}\n            {%- endfor %}\n        {%- else %}\n            {{- message[\"content\"] }}\n        {%- endif %}\n        {{- \"[/INST]\" }}\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 %}"
+}

special_tokens_map.json

@@ -1,5 +1,23 @@
 {
-  "bos_token": "<s>",
-  "eos_token": "</s>",
-  "unk_token": "<unk>"
+  "bos_token": {
+    "content": "<s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "</s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
 }

tokenizer_config.json

@@ -8004,9 +8004,12 @@
   "bos_token": "<s>",
   "clean_up_tokenization_spaces": true,
   "eos_token": "</s>",
+  "model_input_names": [
+    "input_ids",
+    "attention_mask"
+  ],
   "model_max_length": 1000000000000000019884624838656,
   "processor_class": "PixtralProcessor",
   "tokenizer_class": "PreTrainedTokenizerFast",
-  "model_input_names": ["input_ids", "attention_mask"],
   "unk_token": "<unk>"
 }