ColossalAI/applications/Colossal-LLaMA/stream_chat_example.py

61 lines
2.4 KiB
Python
Raw Normal View History

import argparse
from colossal_llama.utils.stream_chat_patch import streaming_chat
from transformers import AutoModelForCausalLM, AutoTokenizer
SYSTEM = "A chat between a curious human and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the human's questions."
def main(args):
model = AutoModelForCausalLM.from_pretrained(args.model_path).cuda().eval()
tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_path)
past_key_values, history = None, []
roles = ["", "Human", "Assistant"]
history = []
history.append({"role": roles[0], "message": SYSTEM})
while True:
input_query = input(f"\n{roles[1]}: ")
if input_query.strip() == "exit":
break
if input_query.strip() == "clear":
past_key_values, history = None, []
continue
print(f"\n{roles[2]}: ", end="")
gen_len = 0
for response, history, past_key_values in streaming_chat(
model,
tokenizer,
input_query,
history=history,
roles=roles,
temperature=args.temperature,
top_p=args.top_p,
top_k=args.top_k,
do_sample=args.do_sample,
length_penalty=args.length_penalty,
max_new_tokens=args.max_new_tokens,
past_key_values=past_key_values,
return_past_key_values=True,
):
output = response[gen_len:]
print(output, end="", flush=True)
gen_len = len(response)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--model_path", type=str, default=None, help="path to chat version model")
parser.add_argument("--tokenizer_path", type=str, default=None, help="path to chat version tokenizer")
parser.add_argument("--temperature", type=float, default=0.8, help="set temperature")
parser.add_argument("--top_p", type=float, default=0.95, help="set top p value")
parser.add_argument("--top_k", type=int, default=50, help="set top k value")
parser.add_argument("--do_sample", type=bool, default=True, help="whether turn on do_sample or not")
parser.add_argument("--length_penalty", type=float, default=1.2, help="set length penalty")
parser.add_argument("--max_new_tokens", type=int, default=512, help="set max new tokens")
args = parser.parse_args()
main(args)