import os
import gradio as gr
import clueai
import torch
from transformers import T5Tokenizer, T5ForConditionalGeneration


base_info = ""




def clear_session():
    return '', None


def ChatYuan(api_key, text_prompt, top_p):
    cl = clueai.Client(api_key, check_api_key=True)
    # generate a prediction for a prompt
    # 需要返回得分的话，指定return_likelihoods="GENERATION"
    prediction = cl.generate(model_name='ChatYuan-7B', prompt=text_prompt)
    # print the predicted text
    #print('prediction: {}'.format(prediction.generations[0].text))
    response = prediction.generations[0].text
    if response == '':
        response = "很抱歉，我无法回答这个问题"

    return response


def chatyuan_bot_api(api_key, input, history, top_p, num):
    if not api_key:
        return "请填写api key再尝试输入", [], []
    history = history or []

    if len(history) > num:
        history = history[-num:]

    #print(f"history{history}")
    history_context = [f"用户：{input_text}\n小元：{answer_text}" for input_text, answer_text in history]
    context = "\n".join(history_context)
    #print(f"context:{context}")
    
    while len(context) > 768:
        history_context = history_context[1:]
        context = "\n".join(history_context)
        
    input_text = context + "\n用户：" + input + "\n小元："
    input_text = input_text.strip()
    output_text = ChatYuan(api_key, input_text, top_p)
    print("api".center(20, "="))
    print(f"api_key:{api_key}\n{input_text}\n{output_text}")

    history.append((input, output_text))

    return '', history, history


block = gr.Blocks()

with block as demo_1:
    gr.Markdown("""<h1><center>元语智能——ChatYuan</center></h1>
    <font size=4>回答来自ChatYuan, 以上是模型生成的结果, 请谨慎辨别和参考, 不代表任何人观点  | Answer generated by ChatYuan model</font>
    <font size=4>注意：gradio对markdown代码格式展示有限</font>
    <font size=4>在使用此功能前，你需要有个API key. API key 可以通过这个<a href='https://www.clueai.cn/' target="_blank">平台</a>获取</font>
    """)
    with gr.Row():
        with gr.Column(scale=3):
            chatbot = gr.Chatbot(label='ChatYuan').style(height=400)

        with gr.Column(scale=1):
            api_key = gr.inputs.Textbox(label="请输入你的api-key(必填)",
                                        default="",
                                        type='password')
            num = gr.Slider(minimum=4,
                            maximum=10,
                            label="最大的对话轮数",
                            value=5,
                            step=1)
            top_p = gr.Slider(minimum=0,
                              maximum=1,
                              label="top_p",
                              value=0.7,
                              step=0.1)
            clear_history = gr.Button("👋 清除历史对话 | Clear History")
            send = gr.Button("🚀 发送 | Send")

    message = gr.Textbox()
    state = gr.State()
    message.submit(chatyuan_bot_api,
                   inputs=[api_key, message, state, top_p, num],
                   outputs=[message, chatbot, state])

    send.click(chatyuan_bot_api,
               inputs=[api_key, message, state, top_p, num],
               outputs=[message, chatbot, state])
    clear_history.click(fn=clear_session,
                        inputs=[],
                        outputs=[chatbot, state],
                        queue=False)

block = gr.Blocks()
with block as introduction:
    gr.Markdown("""<h1><center>元语智能——ChatYuan</center></h1>
    
<font size=4>😉ChatYuan: 元语功能型对话大模型 | General Model for Dialogue with ChatYuan
<br>
👏ChatYuan-7B是一个支持中英双语的功能型对话语言大模型。
<br>
ChatYuan-7B is an open-source large language model for dialogue, supports both Chinese and English languages, and in ChatGPT style.
<br>
ChatYuan-7B是ChatYuan系列中以轻量化实现高质量效果的模型之一。
<br>
<br>
<br>
< br>
</font>
<center><a href="https://clustrmaps.com/site/1bts0"  title="Visit tracker"><img src="//www.clustrmaps.com/map_v2.png?d=ycVCe17noTYFDs30w7AmkFaE-TwabMBukDP1802_Lts&cl=ffffff" /></a></center>
    """)

gui = gr.TabbedInterface(
    interface_list=[introduction, demo_1],
    tab_names=["相关介绍 | Introduction", "API调用"])
gui.launch(quiet=True, show_api=False, share=False)