From b98635d823714818b863d6f82741c705e1223b80 Mon Sep 17 00:00:00 2001
From: pandora <128635000+pandora-s-git@users.noreply.github.com>
Date: Sat, 21 Sep 2024 18:48:28 +0200
Subject: [PATCH] Fixing Mistral Templates

Hi there, fixing the templates as close as possible to the ground truth, you ca find more information regarding the templates in here: https://github.com/mistralai/cookbook/blob/main/concept-deep-dive/tokenization/chat_templates.md

Still needs to be verified, so please dont merge yet!
---
 instruction-templates/Mistral.yaml | 38 +++++++++++++++++++-----------
 1 file changed, 24 insertions(+), 14 deletions(-)

diff --git a/instruction-templates/Mistral.yaml b/instruction-templates/Mistral.yaml
index 02e5b20d..1bb33ded 100644
--- a/instruction-templates/Mistral.yaml
+++ b/instruction-templates/Mistral.yaml
@@ -1,15 +1,25 @@
 instruction_template: |-
-  {%- for message in messages %}
-      {%- if message['role'] == 'system' -%}
-          {{- message['content'] -}}
-      {%- else -%}
-          {%- if message['role'] == 'user' -%}
-              {{-'[INST] ' + message['content'].rstrip() + ' [/INST]'-}}
-          {%- else -%}
-              {{-'' + message['content'] + '</s>' -}}
-          {%- endif -%}
-      {%- endif -%}
-  {%- endfor -%}
-  {%- if add_generation_prompt -%}
-      {{-''-}}
-  {%- endif -%}
\ No newline at end of file
+  {%- if messages[0]['role'] == 'system' %}
+      {%- set system_message = messages[0]['content'] %}
+      {%- set loop_messages = messages[1:] %}
+  {%- else %}
+      {%- set loop_messages = messages %}
+  {%- endif %}
+  
+  {{- bos_token }}
+  {%- for message in loop_messages %}
+      {%- if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}
+          {{- raise_exception('After the optional system message, conversation roles must alternate user/assistant/user/assistant/...') }}
+      {%- endif %}
+      {%- if message['role'] == 'user' %}
+          {%- if loop.first and system_message is defined %}
+              {{- ' [INST] ' + system_message + '\\n\\n' + message['content'] + ' [/INST]' }}
+          {%- else %}
+              {{- ' [INST] ' + message['content'] + ' [/INST]' }}
+          {%- endif %}
+      {%- elif message['role'] == 'assistant' %}
+          {{- ' ' + message['content'] + '</s>'}}
+      {%- else %}
+          {{- raise_exception('Only user and assistant roles are supported, with the exception of an initial optional system message!') }}
+      {%- endif %}
+  {%- endfor %}