From 44b4a0c9c5f65ca8c6d0bca8416be6f86edbf026 Mon Sep 17 00:00:00 2001
From: Ziteng Zhang <87107332+Jasonzzt@users.noreply.github.com>
Date: Tue, 26 Dec 2023 16:57:55 +0800
Subject: [PATCH] [LLM] Correct prompt format of Yi, Llama2 and Qwen in
generate.py (#9786)
* correct prompt format of Yi
* correct prompt format of llama2 in cpu generate.py
* correct prompt format of Qwen in GPU example
---
.../Model/llama2/generate.py | 11 ++++++-----
.../HF-Transformers-AutoModels/Model/yi/generate.py | 10 ++++++++--
.../HF-Transformers-AutoModels/Model/qwen/generate.py | 10 +++++++++-
.../HF-Transformers-AutoModels/Model/yi/generate.py | 10 ++++++++--
4 files changed, 31 insertions(+), 10 deletions(-)
diff --git a/python/llm/example/CPU/HF-Transformers-AutoModels/Model/llama2/generate.py b/python/llm/example/CPU/HF-Transformers-AutoModels/Model/llama2/generate.py
index 6353d848..1d6c56e2 100644
--- a/python/llm/example/CPU/HF-Transformers-AutoModels/Model/llama2/generate.py
+++ b/python/llm/example/CPU/HF-Transformers-AutoModels/Model/llama2/generate.py
@@ -22,11 +22,12 @@ from bigdl.llm.transformers import AutoModelForCausalLM
from transformers import LlamaTokenizer
# you could tune the prompt based on your own model,
-# here the prompt tuning refers to https://huggingface.co/georgesung/llama2_7b_chat_uncensored#prompt-style
-LLAMA2_PROMPT_FORMAT = """### HUMAN:
-{prompt}
-
-### RESPONSE:
+# Refer to https://huggingface.co/TheBloke/Llama-2-70B-Chat-GGML#prompt-template-llama-2-chat
+LLAMA2_PROMPT_FORMAT = """
+[INST] <>
+You are a helpful assistant.
+<>
+{prompt}[/INST]
"""
if __name__ == '__main__':
diff --git a/python/llm/example/CPU/HF-Transformers-AutoModels/Model/yi/generate.py b/python/llm/example/CPU/HF-Transformers-AutoModels/Model/yi/generate.py
index 161795d6..c7024b3c 100644
--- a/python/llm/example/CPU/HF-Transformers-AutoModels/Model/yi/generate.py
+++ b/python/llm/example/CPU/HF-Transformers-AutoModels/Model/yi/generate.py
@@ -21,8 +21,14 @@ import argparse
from bigdl.llm.transformers import AutoModelForCausalLM
from transformers import AutoTokenizer
-# you could tune the prompt based on your own model,
-YI_PROMPT_FORMAT = "{prompt}"
+# Refer to https://huggingface.co/01-ai/Yi-6B-Chat#31-use-the-chat-model
+YI_PROMPT_FORMAT = """
+<|im_start|>system
+You are a helpful assistant. If you don't understand what the user means, ask the user to provide more information.<|im_end|>
+<|im_start|>user
+{prompt}<|im_end|>
+<|im_start|>assistant
+"""
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='Predict Tokens using `generate()` API for Yi model')
diff --git a/python/llm/example/GPU/HF-Transformers-AutoModels/Model/qwen/generate.py b/python/llm/example/GPU/HF-Transformers-AutoModels/Model/qwen/generate.py
index 9e998470..0b898cde 100644
--- a/python/llm/example/GPU/HF-Transformers-AutoModels/Model/qwen/generate.py
+++ b/python/llm/example/GPU/HF-Transformers-AutoModels/Model/qwen/generate.py
@@ -23,7 +23,15 @@ from bigdl.llm.transformers import AutoModelForCausalLM
from transformers import AutoTokenizer
# you could tune the prompt based on your own model
-QWEN_PROMPT_FORMAT = "{prompt} "
+QWEN_PROMPT_FORMAT = """
+<|im_start|>system
+You are a helpful assistant.
+<|im_end|>
+<|im_start|>user
+{prompt}
+<|im_end|>
+<|im_start|>assistant
+"""
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='Predict Tokens using `generate()` API for Qwen model')
diff --git a/python/llm/example/GPU/HF-Transformers-AutoModels/Model/yi/generate.py b/python/llm/example/GPU/HF-Transformers-AutoModels/Model/yi/generate.py
index 13a9f61f..214e2904 100644
--- a/python/llm/example/GPU/HF-Transformers-AutoModels/Model/yi/generate.py
+++ b/python/llm/example/GPU/HF-Transformers-AutoModels/Model/yi/generate.py
@@ -22,8 +22,14 @@ import argparse
from bigdl.llm.transformers import AutoModelForCausalLM
from transformers import AutoTokenizer
-# you could tune the prompt based on your own model
-YI_PROMPT_FORMAT = "{prompt}"
+# Refer to https://huggingface.co/01-ai/Yi-6B-Chat#31-use-the-chat-model
+YI_PROMPT_FORMAT = """
+<|im_start|>system
+You are a helpful assistant. If you don't understand what the user means, ask the user to provide more information.<|im_end|>
+<|im_start|>user
+{prompt}<|im_end|>
+<|im_start|>assistant
+"""
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='Predict Tokens using `generate()` API for Yi model')