[LLM] Correct prompt format of Yi, Llama2 and Qwen in generate.py (in…

…tel#9786) * correct prompt format of Yi * correct prompt format of llama2 in cpu generate.py * correct prompt format of Qwen in GPU example
WeiguangHan · Dec 26, 2023 · 30a3ba6 · 30a3ba6
1 parent 981731b
commit 30a3ba6
Show file tree

Hide file tree

Showing 4 changed files with 31 additions and 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
@@ -22,11 +22,12 @@
 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] <<SYS>>
+You are a helpful assistant.
+<</SYS>>
+{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
@@ -21,8 +21,14 @@
 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
@@ -23,7 +23,15 @@
 from transformers import AutoTokenizer
 
 # you could tune the prompt based on your own model
-QWEN_PROMPT_FORMAT = "<human>{prompt} <bot>"
+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
@@ -22,8 +22,14 @@
 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')