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")

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