import gradio as gr




import torch
from transformers import AutoTokenizer
from transformers import T5Tokenizer, T5ForConditionalGeneration

# tokenizer = T5Tokenizer.from_pretrained("ClueAI/PromptCLUE-base")
# model = T5ForConditionalGeneration.from_pretrained("ClueAI/PromptCLUE-base")
# tokenizer = T5Tokenizer.from_pretrained("ClueAI/PromptCLUE-base-v1-5")
# model = T5ForConditionalGeneration.from_pretrained("ClueAI/PromptCLUE-base-v1-5")



tokenizer = T5Tokenizer.from_pretrained("ClueAI/ChatYuan-large-v1")
model = T5ForConditionalGeneration.from_pretrained("ClueAI/ChatYuan-large-v1")

device = torch.device('cpu')
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])


#iface = gr.Interface(fn=answer, inputs="text", outputs="text")
examples = [
    ["""摘要这段话: 
     现在乌军的迫击炮都可以开始轰击库皮扬斯克的俄军目标了，双方相距只有几公里。
    最重要的是该地区俄军背后就是奥斯科尔河-北顿涅茨河。而为了确保第聂伯河上赫尔松城的后勤保障，俄军已经把舟桥部队主力调到赫尔松去了，现在是远水解不了近渴。
    乌军很有可能暂时不会直接攻击伊久姆城区，现在还是要先拿下库皮扬斯克和奥斯科尔河上的两座桥梁。
    切断俄军从别尔哥罗德向库皮扬斯克的物资输送，切断库皮扬斯克-伊久姆公路，断其粮道。
    现在要看俄军有没有预备队。库皮扬斯克没援军的话大概率守不住。但是现在打了三天，俄军还没有援军抵达战场。
    现在，俄军最主要任务是守住库皮扬斯克，同时要确保库皮扬斯克-伊久姆高速公路的安全。
    """],
    ["""翻译这段话到英文:
     在北京冬奥会自由式滑雪女子坡面障碍技巧决赛中，中国选手谷爱凌夺得银牌。祝贺谷爱凌！今天上午，自由式滑雪女子坡面障碍技巧决赛举行。决赛分三轮进行，取选手最佳成绩排名决出奖牌。第一跳，中国选手谷爱凌获得69.90分。在12位选手中排名第三。完成动作后，谷爱凌又扮了个鬼脸，甚是可爱。第二轮中，谷爱凌在道具区第三个障碍处失误，落地时摔倒。获得16.98分。网友：摔倒了也没关系，继续加油！在第二跳失误摔倒的情况下，谷爱凌顶住压力，第三跳稳稳发挥，流畅落地！获得86.23分！此轮比赛，共12位选手参赛，谷爱凌第10位出场。网友：看比赛时我比谷爱凌紧张，加油！
     """],
]

iface = gr.Interface(
    fn=answer,
    inputs=gr.Textbox(lines=5, label="Input Text"),
    outputs=gr.Textbox(label="Generated Text"),
    examples=examples
)
iface.launch()