# 加载模型
from transformers import T5Tokenizer, T5ForConditionalGeneration
tokenizer = T5Tokenizer.from_pretrained("ClueAI/ChatYuan-large-v1")
model = T5ForConditionalGeneration.from_pretrained("ClueAI/ChatYuan-large-v1")

# 使用模型进行预测推理方法：
import torch
from transformers import AutoTokenizer
# 修改colab笔记本设置为gpu，推理更快
device = torch.device('cuda')
model.to(device)

def preprocess(text):  
	text = text.replace("\n", "\\n").replace("\t", "\\t")  
	return text

def postprocess(text):  
	return text.replace("\\n", "\n").replace("\\t", "\t")
	
def answer(text, sample=True, top_p=1, temperature=0.7):  
	'''sample：是否抽样。生成任务，可以设置为True;  
	top_p：0-1之间，生成的内容越多样'''  
	text = preprocess(text)  
	encoding = tokenizer(text=[text], truncation=True, padding=True, max_length=768, 		return_tensors="pt").to(device)   
	if not sample:    
		out = model.generate(**encoding, return_dict_in_generate=True, output_scores=False, max_new_tokens=512, num_beams=1, length_penalty=0.6)  
	else:    
		out = model.generate(**encoding, return_dict_in_generate=True, output_scores=False, max_new_tokens=512, do_sample=True, top_p=top_p, temperature=temperature, no_repeat_ngram_size=3)  
	out_text = tokenizer.batch_decode(out["sequences"], skip_special_tokens=True)  
	return postprocess(out_text[0])
print("end...")

input_text0 = "帮我写一个请假条，我因为新冠不舒服，需要请假3天，请领导批准"
input_text1 = "你能干什么"
input_text2 = "用英文写一封道歉的邮件，表达因为物流延误，不能如期到达，我们可以赔偿贵公司所有损失"
input_text3 = "写一个文章，题目是未来城市"
input_text4 = "写一个诗歌，关于冬天"
input_text5 = "从南京到上海的路线"
input_text6 = "学前教育专业岗位实习中，在学生方面会存在问题，请提出改进措施。800字"
input_text7 = "根据标题生成文章：标题：屈臣氏里的化妆品到底怎么样？正文：化妆品，要讲究科学运用，合理搭配。屈臣氏起码是正品连锁店。请继续后面的文字。"
input_text8 = "帮我对比几款GPU，列出详细参数对比，并且给出最终结论"
input_list = [input_text0, input_text1, input_text2, input_text3, input_text4, input_text5, input_text6, input_text7, input_text8]
for i, input_text in enumerate(input_list):  
	input_text = "用户：" + input_text + "\n小元："  
	print(f"示例{i}".center(50, "="))  
	output_text = answer(input_text)  
	print(f"{input_text}{output_text}")

input_text = ["你好","新冠什么症状？","可以吃什么药？"]
answer_text = ["您好!我是元语AI。我可以回答您的问题、写文章、写作业、翻译，对于一些法律等领域的问题我也可以给你提供信息", "新冠是指新型冠状病毒，其症状包括发热、干咳、乏力、嗅味觉减退、呼吸困难等。", "根据您提供的病史，目前没有明确的抗新冠病毒的药物，建议您在家进行自我隔离，避免与他人接触，多喝开水，清淡易消化饮食，避免熬夜和过度劳累，适当进行户外活动。"]
context = "\n".join([f"用户：{input_text[i]}\n小元：{answer_text[i]}" for i in range(len(input_text))])
print(context)

input_text = "用什么后遗症么？"
print(f"示例".center(50, "="))
input_text = context + "\n用户：" + input_text + "\n小元："
output_text = answer(input_text)
print(f"{input_text}{output_text}")

# 加载模型
from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks
from modelscope.models.nlp import T5ForConditionalGeneration
from modelscope.preprocessors import Text2TextGenerationPreprocessor

# 使用模型进行推理：
model = T5ForConditionalGeneration.from_pretrained('ClueAI/ChatYuan-large', revision='v1.0.0')
preprocessor = Text2TextGenerationPreprocessor(model.model_dir)
pipeline_t2t = pipeline(task=Tasks.text2text_generation, model=model, preprocessor=preprocessor)  
# 生成
print(pipeline_t2t('用户：帮我写个请假条，我因为新冠不舒服，需要请假3天，请领导批准\\n小元：', do_sample=True, top_p=0.9, max_length=256))'''

尊敬的领导：
您好!
因为最近新冠病毒的流行，我感到身体不舒服，所以请求请假3天，在这段时间内，我将积极配合医生的治疗，保证充足的睡眠，并做好防护工作。
在此期间，我保证将努力完成我的工作，并按时下班，在请假期间，我绝对遵守公司的规章制度，接受公司的监督和检查，请领导批准。谢谢!_ 
此致
敬礼!'''    
# 多轮对话print(pipeline_t2t('用户：新冠什么症状？\\n小元：新冠是指新型冠状病毒，其症状包括发热、干咳、乏力、嗅味觉减退、呼吸困难等。\\n用户：可以吃什么药？\\n小元：根据您提供的病史，目前没有明确的抗新冠病毒的药物，建议您在家进行自我隔离，避免与他人接触，多喝开水，清淡易消化饮食，避免熬夜和过度劳累，适当进行户外活动。\\n用户：用什么后遗症么？\\n小元：', do_sample=True, top_p=0.9, max_length=256))'''目前还没有人具体说是什么后遗症，但是目前症状比较轻的，可能没有后遗症，但是如果症状比较重，就可能出现呼吸困难，胸闷，发热，咳嗽等症状。'''========================示例========================  用户：新冠什么症状？小元：新冠是指新型冠状病毒，其症状包括发热、干咳、乏力、嗅味觉减退、呼吸困难等。用户：可以吃什么药？小元：根据您提供的病史，目前没有明确的抗新冠病毒的药物，建议您在家进行自我隔离，避免与他人接触，多喝开水，清淡易消化饮食，避免熬夜和过度劳累，适当进行户外活动。用户：用什么后遗症么？----小元：目前还没有人具体说是什么后遗症，但是目前症状比较轻的，可能没有后遗症，但是如果症状比较重，就可能出现呼吸困难，胸闷，发热，咳嗽等症状。