import torch
import transformers


def apply_multi_turn_template(history, input):
    multi_turn_template = "[ROUND {} USER]{}[ROUND {} ASSISTANT]{}"
    prefix = ""
    for i in range(len(history)):
        prefix = prefix + multi_turn_template.format(i, history[i][0], i, history[i][1])
        prefix = prefix + "<|end_of_turn|>"
    curr_turn_id = len(history)
    curr_turn = "[ROUND {} USER]{}[ROUND {} ASSISTANT]".format(curr_turn_id, input, curr_turn_id)
    return prefix + curr_turn
    

model = transformers.AutoModelForCausalLM.from_pretrained("XuYipei/Mixtral-4x7B-Chat-Chinese")
tokenizer = transformers.AutoTokenizer.from_pretrained("XuYipei/Mixtral-4x7B-Chat-Chinese")

model = model.to(torch.bfloat16).to("cuda")
input_text = apply_multi_turn_template([], "列举出10首杜甫的古诗，同时默写该诗篇的完整内容。")
output = md.generate(input_ids = torch.LongTensor([tk.encode(input_text)]).to("cuda:5"), eos_token_id=55296, max_length=4096)
output_text = tokenizer.decode(output[0])

'''
<s> [ROUND 0 USER]列举出10首杜甫的古诗，同时默写该诗篇的完整内容。[ROUND 0 ASSISTANT] 1. 《登高》

唐·杜甫

风急天高猿啸哀，渚清沙白鸟飞回。
无边落木萧萧下，不尽长江滚滚来。
万里悲秋常作客，百年多病独登台。
艰难苦恨繁霜鬓，潦倒新停浊酒杯。

2. 《春夜喜雨》

唐·杜甫

好雨知时节，当春乃发生。
随风潜入夜，润物细无声。
野径云俱黑，江船火独明。
晓看红湿处，花重锦官城。

3. 《绝句》

唐·杜甫

迟日江山丽，春风花草香。
泥融飞燕子，沙暖睡鸳鸯。

4. 《登楼》

唐·杜甫

花近高楼伤客心，万方多难此登临。
锦江春色来天地，玉垒浮云变古今。
北极朝廷终不改，西山寇盗莫相侵。
可怜后主还祠庙，日暮聊为梁甫吟。

5. 《茅屋为秋风所破歌》

唐·杜甫

八月秋高风怒号，卷我屋上三重茅。
茅飞渡江洒江郊，高者挂罥长林梢，
下者飘转沉塘坳。
南村群童欺我老无力，忍能对面为盗贼。
公然抱茅入竹去，唇焦口燥呼不得，
归来倚杖自叹息。

6. 《江南逢李龟年》

唐·杜甫

岐王宅里寻常见，崔九堂前几度闻。
正是江南好风景，落花时节又逢君。

7. 《闻官军收河南河北》

唐·杜甫

剑外忽传收蓟北，初闻涕泪满衣裳。
却看妻子愁何在，漫卷诗书喜欲狂。
白日放歌须纵酒，青春作伴好还乡。
即从巴峡穿巫峡，便下襄阳向洛阳。

8. 《旅夜书怀》

唐·杜甫

细草微风岸，危樯独夜舟。
星垂平野阔，月涌大江流。
名岂文章著，官应老病休。
飘飘何所似，天地一沙鸥。

9. 《江畔独步寻花》

唐·杜甫

黄四娘家花满蹊，千朵万朵压枝低。
留恋戏蝶时时舞，自在娇莺恰恰啼。

10. 《赠花卿》

唐·杜甫

锦城丝管日纷纷，半入江风半入云。
此曲只应天上有，人间能得几回闻。<|end_of_turn|>

'''

致谢

感谢知识工厂梁家卿老师、肖仰华老师的指导，感谢小i机器人的支持。