import os
# install torch and tf
os.system('pip install transformers SentencePiece')
os.system('pip install torch')

# pip install streamlit-chat 
os.system('pip install streamlit-chat')

from transformers import T5Tokenizer, T5ForConditionalGeneration, AutoTokenizer
import torch

import streamlit as st
from streamlit_chat import message

# 下载模型
tokenizer = T5Tokenizer.from_pretrained("ClueAI/ChatYuan-large-v1")
model = T5ForConditionalGeneration.from_pretrained("ClueAI/ChatYuan-large-v1")
# 修改colab笔记本设置为gpu，推理更快
device = torch.device('cpu')
model.to(device)
print('Model Load done!')

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(user_history, bot_history, sample=True, top_p=1, temperature=0.7):
    '''sample：是否抽样。生成任务，可以设置为True;
    top_p：0-1之间，生成的内容越多样
    max_new_tokens=512 lost...'''

    if len(bot_history)>0:
        context = "\n".join([f"用户：{user_history[i]}\n小元：{bot_history[i]}" for i in range(len(bot_history))])
        input_text = context + "\n用户：" + user_history[-1] + "\n小元："
    else:
        input_text = "用户：" + user_history[-1] + "\n小元："
    
    input_text = preprocess(input_text)
    print(input_text)
    encoding = tokenizer(text=input_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)
    print('小元: '+postprocess(out_text[0]))
    return postprocess(out_text[0])

st.set_page_config(
    page_title="Chinese ChatBot - Demo",
    page_icon=":robot:"
)

st.header("Chinese ChatBot - Demo")
st.markdown("[Github](https://github.com/scutcyr)")

if 'generated' not in st.session_state:
    st.session_state['generated'] = []

if 'past' not in st.session_state:
    st.session_state['past'] = []


def get_text():
    input_text = st.text_input("用户: ","你好！", key="input")
    return input_text  

user_history = []
bot_history = []
user_input = get_text()
user_history.append(user_input)

if user_input:
    output = answer(user_history,bot_history)
    st.session_state.past.append(user_input)
    st.session_state.generated.append(output)
    bot_history.append(output)

if st.session_state['generated']:

    #for i in range(len(st.session_state['generated'])-1, -1, -1):
    #    message(st.session_state["generated"][i], key=str(i))
    #    message(st.session_state['past'][i], is_user=True, key=str(i) + '_user')
    for i in range(len(st.session_state['generated'])):
        message(st.session_state['past'][i], is_user=True, key=str(i) + '_user')
        message(st.session_state["generated"][i], key=str(i))