From 3186bd7e34da5cdd15a03f685185dadf478d50fe Mon Sep 17 00:00:00 2001
From: Josh <josh@civo.com>
Date: Mon, 10 Jun 2024 16:27:29 +0100
Subject: [PATCH] including chat templates

---
 Dockerfile                           |  3 +++
 chat-templates/alpacha.jinja         | 29 ++++++++++++++++++++++++++++
 chat-templates/chatglm.jinja         | 18 +++++++++++++++++
 chat-templates/chatglm2.jinja        | 18 +++++++++++++++++
 chat-templates/chatml.jinja          |  2 ++
 chat-templates/falcon-180b.jinja     | 17 ++++++++++++++++
 chat-templates/falcon.jinja          | 15 ++++++++++++++
 chat-templates/llama2.jinja          | 24 +++++++++++++++++++++++
 chat-templates/llama3-instruct.jinja | 18 +++++++++++++++++
 chat-templates/llava.jinja           | 23 ++++++++++++++++++++++
 entrypoint.sh                        |  4 ++++
 requirements.txt                     |  2 +-
 12 files changed, 172 insertions(+), 1 deletion(-)
 create mode 100644 chat-templates/alpacha.jinja
 create mode 100644 chat-templates/chatglm.jinja
 create mode 100644 chat-templates/chatglm2.jinja
 create mode 100644 chat-templates/chatml.jinja
 create mode 100644 chat-templates/falcon-180b.jinja
 create mode 100644 chat-templates/falcon.jinja
 create mode 100644 chat-templates/llama2.jinja
 create mode 100644 chat-templates/llama3-instruct.jinja
 create mode 100644 chat-templates/llava.jinja

diff --git a/Dockerfile b/Dockerfile
index d1c4ae9..fc4611a 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -10,6 +10,9 @@ COPY --chown=$MAMBA_USER:$MAMBA_USER requirements.txt /tmp/requirements.txt
 # install the python deps
 RUN python3 -m pip install --no-cache-dir -r /tmp/requirements.txt
 
+# copy over the chat templates
+COPY chat-templates /chat-templates
+
 # copy over the entry point
 COPY --chown=$MAMBA_USER:$MAMBA_USER entrypoint.sh /
 RUN chmod +x /entrypoint.sh
diff --git a/chat-templates/alpacha.jinja b/chat-templates/alpacha.jinja
new file mode 100644
index 0000000..60667ac
--- /dev/null
+++ b/chat-templates/alpacha.jinja
@@ -0,0 +1,29 @@
+{{ (messages|selectattr('role', 'equalto', 'system')|list|last).content|trim if (messages|selectattr('role', 'equalto', 'system')|list) else '' }}
+
+{% for message in messages %}
+{% if message['role'] == 'user' %}
+### Instruction:
+{{ message['content']|trim -}}
+{% if not loop.last %}
+
+
+{% endif %}
+{% elif message['role'] == 'assistant' %}
+### Response:
+{{ message['content']|trim -}}
+{% if not loop.last %}
+
+
+{% endif %}
+{% elif message['role'] == 'user_context' %}
+### Input:
+{{ message['content']|trim -}}
+{% if not loop.last %}
+
+
+{% endif %}
+{% endif %}
+{% endfor %}
+{% if add_generation_prompt and messages[-1]['role'] != 'assistant' %}
+### Response:
+{% endif %}
\ No newline at end of file
diff --git a/chat-templates/chatglm.jinja b/chat-templates/chatglm.jinja
new file mode 100644
index 0000000..bf26f27
--- /dev/null
+++ b/chat-templates/chatglm.jinja
@@ -0,0 +1,18 @@
+{%- set counter = namespace(index=0) -%}
+{%- for message in messages -%}
+    {%- if message['role'] == 'user' -%}
+        {{- '[Round ' + counter.index|string + ']\n问：' + message['content'] -}}
+        {%- set counter.index = counter.index + 1 -%}
+    {%- endif -%}
+    {%- if message['role'] == 'assistant' -%}
+        {{- '\n答：' + message['content'] -}}
+        {%- if (loop.last and add_generation_prompt) or not loop.last -%}
+            {{- '\n' -}}
+        {%- endif -%}
+    {%- endif -%}
+{%- endfor -%}
+
+
+{%- if add_generation_prompt and messages[-1]['role'] != 'assistant' -%}
+    {{- '\n答：' -}}
+{%- endif -%}
\ No newline at end of file
diff --git a/chat-templates/chatglm2.jinja b/chat-templates/chatglm2.jinja
new file mode 100644
index 0000000..c155b7c
--- /dev/null
+++ b/chat-templates/chatglm2.jinja
@@ -0,0 +1,18 @@
+{%- set counter = namespace(index=1) -%}
+{%- for message in messages -%}
+    {%- if message['role'] == 'user' -%}
+        {{- '[Round ' + counter.index|string + ']\n\n问：' + message['content'] -}}
+        {%- set counter.index = counter.index + 1 -%}
+    {%- endif -%}
+    {%- if message['role'] == 'assistant' -%}
+        {{- '\n\n答：' + message['content'] -}}
+        {%- if (loop.last and add_generation_prompt) or not loop.last -%}
+            {{- '\n\n' -}}
+        {%- endif -%}
+    {%- endif -%}
+{%- endfor -%}
+
+
+{%- if add_generation_prompt and messages[-1]['role'] != 'assistant' -%}
+    {{- '\n\n答：' -}}
+{%- endif -%}
\ No newline at end of file
diff --git a/chat-templates/chatml.jinja b/chat-templates/chatml.jinja
new file mode 100644
index 0000000..4844e68
--- /dev/null
+++ b/chat-templates/chatml.jinja
@@ -0,0 +1,2 @@
+{% for message in messages %}{{'<|im_start|>' + message['role'] + '\n' + message['content']}}{% if (loop.last and add_generation_prompt) or not loop.last %}{{ '<|im_end|>' + '\n'}}{% endif %}{% endfor %}
+{% if add_generation_prompt and messages[-1]['role'] != 'assistant' %}{{ '<|im_start|>assistant\n' }}{% endif %}
\ No newline at end of file
diff --git a/chat-templates/falcon-180b.jinja b/chat-templates/falcon-180b.jinja
new file mode 100644
index 0000000..f08f739
--- /dev/null
+++ b/chat-templates/falcon-180b.jinja
@@ -0,0 +1,17 @@
+{%- for message in messages -%}
+    {%- if message['role'] == 'system' -%}
+        {{- 'System: ' + message['content'] -}}
+    {%- elif message['role'] == 'user' -%}
+        {{- 'User: ' + message['content'] -}}
+    {%- elif message['role'] == 'assistant' -%}
+        {{- 'Falcon: ' + message['content'] -}}
+    {%- endif -%}
+    {%- if (loop.last and add_generation_prompt) or not loop.last -%}
+        {{- '\n' -}}
+    {%- endif -%}
+{%- endfor -%}
+
+
+{%- if add_generation_prompt and messages[-1]['role'] != 'assistant' -%}
+    {{- 'Falcon:' -}}
+{% endif %}
\ No newline at end of file
diff --git a/chat-templates/falcon.jinja b/chat-templates/falcon.jinja
new file mode 100644
index 0000000..01cf0e2
--- /dev/null
+++ b/chat-templates/falcon.jinja
@@ -0,0 +1,15 @@
+{%- for message in messages -%}
+    {%- if message['role'] == 'user' -%}
+        {{- 'User: ' + message['content'] -}}
+    {%- elif message['role'] == 'assistant' -%}
+        {{- 'Assistant: ' + message['content'] -}}
+    {%- endif -%}
+    {%- if (loop.last and add_generation_prompt) or not loop.last -%}
+        {{- '\n' -}}
+    {%- endif -%}
+{%- endfor -%}
+
+
+{%- if add_generation_prompt and messages[-1]['role'] != 'assistant' -%}
+    {{- 'Assistant:' -}}
+{% endif %}
\ No newline at end of file
diff --git a/chat-templates/llama2.jinja b/chat-templates/llama2.jinja
new file mode 100644
index 0000000..8b0bc1e
--- /dev/null
+++ b/chat-templates/llama2.jinja
@@ -0,0 +1,24 @@
+{% if messages[0]['role'] == 'system' %}
+    {% set system_message = '<<SYS>>\n' + messages[0]['content'] | trim + '\n<</SYS>>\n\n' %}
+    {% set messages = messages[1:] %}
+{% else %}
+    {% set system_message = '' %}
+{% endif %}
+
+{% for message in messages %}
+    {% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}
+        {{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}
+    {% endif %}
+
+    {% if loop.index0 == 0 %}
+        {% set content = system_message + message['content'] %}
+    {% else %}
+        {% set content = message['content'] %}
+    {% endif %}
+
+    {% if message['role'] == 'user' %}
+        {{ bos_token + '[INST] ' + content | trim + ' [/INST]' }}
+    {% elif message['role'] == 'assistant' %}
+        {{ ' ' + content | trim + ' ' + eos_token }}
+    {% endif %}
+{% endfor %}
\ No newline at end of file
diff --git a/chat-templates/llama3-instruct.jinja b/chat-templates/llama3-instruct.jinja
new file mode 100644
index 0000000..40dbc41
--- /dev/null
+++ b/chat-templates/llama3-instruct.jinja
@@ -0,0 +1,18 @@
+{% if messages[0]['role'] == 'system' %}
+    {% set offset = 1 %}
+{% else %}
+    {% set offset = 0 %}
+{% endif %}
+
+{{ bos_token }}
+{% for message in messages %}
+    {% if (message['role'] == 'user') != (loop.index0 % 2 == offset) %}
+        {{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}
+    {% endif %}
+
+    {{ '<|start_header_id|>' + message['role'] + '<|end_header_id|>\n\n' + message['content'] | trim + '<|eot_id|>' }}
+{% endfor %}
+
+{% if add_generation_prompt %}
+    {{ '<|start_header_id|>' + 'assistant' + '<|end_header_id|>\n\n' }}
+{% endif %}
\ No newline at end of file
diff --git a/chat-templates/llava.jinja b/chat-templates/llava.jinja
new file mode 100644
index 0000000..af8d0b8
--- /dev/null
+++ b/chat-templates/llava.jinja
@@ -0,0 +1,23 @@
+{%- if messages[0]['role'] == 'system' -%}
+    {%- set system_message = messages[0]['content'] -%}
+    {%- set messages = messages[1:] -%}
+{%- else -%}
+    {% set system_message = '' -%}
+{%- endif -%}
+
+{{ bos_token + system_message }}
+{%- 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' -%}
+        {{ 'USER: ' + message['content'] + '\n' }}
+    {%- elif message['role'] == 'assistant' -%}
+        {{ 'ASSISTANT: ' + message['content'] + eos_token + '\n' }}
+    {%- endif -%}
+{%- endfor -%}
+
+{%- if add_generation_prompt -%}
+    {{ 'ASSISTANT:' }}
+{% endif %}
\ No newline at end of file
diff --git a/entrypoint.sh b/entrypoint.sh
index bcaaf7d..0318117 100644
--- a/entrypoint.sh
+++ b/entrypoint.sh
@@ -52,6 +52,10 @@ if [[ -n "$MAX_MODEL_LEN" ]]; then
     additional_args="$additional_args --max-model-len $MAX_MODEL_LEN"
 fi
 
+if [[ ! -z "${CHAT_TEMPLATE}" ]]; then
+    additional_args="${additional_args} --chat-template ${CHAT_TEMPLATE}"
+fi
+
 # Start the API server with the specified or default settings
 python3 -m vllm.entrypoints.openai.api_server \
     --tensor-parallel-size $NUM_GPU \
diff --git a/requirements.txt b/requirements.txt
index 7984ab0..28a4762 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,3 +1,3 @@
-vllm==0.4.2
+vllm==0.4.3
 hf_transfer
 accelerate
\ No newline at end of file