import os

import gradio as gr



os.system(
    'pip install --upgrade torch'
)
os.system(
    'pip install "modelscope" --upgrade -f https://pypi.org/project/modelscope/'
)

os.system('pip install transformers -U')

from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks

pipe = pipeline(task=Tasks.chat, model='ZhipuAI/chatglm2-6b', model_revision='v1.0.2',device_map='auto')


def clear_session():
    return '', None


def predict(input, history=None):
    dictionary  = {
        'prmpt': input
    }
    print(dictionary)
    if history is None:
        history = []
    inputs = inputs = {'text': input, 'history': history}
    result = pipe(inputs)
    history = result['history']

    return '', history, history


block = gr.Blocks()

with block as demo:
    gr.Markdown("""<h1><center>ChatGLM-6B unofficial demo</center></h1>
     <center><font size=3><a href='https://modelscope.cn/models/ZhipuAI/ChatGLM-6B/summary' target="_blank">ChatGLM-6B </a>是一个开源的、支持中英双语的对话语言模型，基于 General Language Model (GLM) 架构，具有 62 亿参数。</center></font>
    """)
    chatbot = gr.Chatbot(label='ChatGLM-6B')
    message = gr.Textbox()
    state = gr.State()
    message.submit(predict,
                   inputs=[message, state],
                   outputs=[message, chatbot, state])
    with gr.Row():
        clear_history = gr.Button("🧹 清除历史对话")
        send = gr.Button("🚀 发送")

    send.click(predict,
               inputs=[message, state],
               outputs=[message, chatbot, state])
    clear_history.click(fn=clear_session,
                        inputs=[],
                        outputs=[chatbot, state],
                        queue=False)

demo.queue().launch(height=800, share=False)