From cb31998605f1887890629f1854d9b147b1b742e4 Mon Sep 17 00:00:00 2001
From: oobabooga <112222186+oobabooga@users.noreply.github.com>
Date: Fri, 3 May 2024 08:19:04 -0700
Subject: [PATCH] Add a template for NVIDIA ChatQA models

---
 instruction-templates/NVIDIA-ChatQA.yaml | 25 ++++++++++++++++++++++++
 models/config.yaml                       |  2 ++
 2 files changed, 27 insertions(+)
 create mode 100644 instruction-templates/NVIDIA-ChatQA.yaml

diff --git a/instruction-templates/NVIDIA-ChatQA.yaml b/instruction-templates/NVIDIA-ChatQA.yaml
new file mode 100644
index 00000000..85a6266b
--- /dev/null
+++ b/instruction-templates/NVIDIA-ChatQA.yaml
@@ -0,0 +1,25 @@
+instruction_template: |-
+  {%- set ns = namespace(found=false) -%}
+  {%- for message in messages -%}
+      {%- if message['role'] == 'system' -%}
+          {%- set ns.found = true -%}
+      {%- endif -%}
+  {%- endfor -%}
+  {%- if not ns.found -%}
+      {{- '' -}}
+  {%- endif %}
+  {%- for message in messages %}
+      {%- if message['role'] == 'system' -%}
+          {{- 'System:' + message['content'] + '\n\n' -}}
+      {%- else -%}
+          {%- if message['role'] == 'user' -%}
+              {{-'User: ' + message['content'] + '\n\n'-}}
+          {%- else -%}
+              {{-'Assistant: ' + message['content'] + '\n\n' -}}
+          {%- endif -%}
+      {%- endif -%}
+  {%- endfor -%}
+  {%- if add_generation_prompt -%}
+      {{-'Assistant:'-}}
+  {%- endif -%}
+
diff --git a/models/config.yaml b/models/config.yaml
index dcbce9e7..8521c4c6 100644
--- a/models/config.yaml
+++ b/models/config.yaml
@@ -204,3 +204,5 @@
   instruction_template: 'ChatML'
 .*airoboros-3_1-yi-34b-200k:
   instruction_template: 'Llama-v2'
+.*chatqa:
+  instruction_template: 'NVIDIA-ChatQA'