from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("svjack/bloom-dialogue")

model = AutoModelForCausalLM.from_pretrained("svjack/bloom-dialogue")

tokenizer.decode(
model.generate(
    tokenizer.encode(
            "你认为今天的天气怎么样？", return_tensors="pt", add_special_tokens=True
        ), max_length = 128,
    num_beams=2,
    top_p = 0.95,
    top_k = 50,
    repetition_penalty = 2.5,
    length_penalty=1.0,
    early_stopping=True,
)[0],
skip_special_tokens = True
).split("\n-----\n")

'''
['你认为今天的天气怎么样？ ',
 ' 下午下雨了，不是吗？ ',
 ' 是的 。 我希望这个周末不会下雨。 ',
 ' 你为什么这么说 ？ ',
 ' 好吧，我们本周要下雨了。 ',
 ' 真的吗 ？ ',
 ' 是的 。 ',
 ' 明天天气会很愉快。 顺便说一句，您想在将来做什么？ ',
 ' 我想成为一名工程师。 ',
 ' 那不是那么容易吗？ ',
 ' 是的，我想成为一名工程师。 ',
 ' 您想要什么样的工作？ ',
 ' 我需要学习成为一名工程师。 ',
 ' 您什么时候开始学习？']
'''