# https://github.com/THUDM/ChatGLM2-6B

from transformers import AutoTokenizer, AutoModel
import gradio as gr

tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True)
# model = AutoModel.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True, device='cuda')
# model = AutoModel.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True).float()
model = AutoModel.from_pretrained("THUDM/chatglm2-6b-int4",trust_remote_code=True).float()
model = model.eval()

def text_generation(prompt: str, style: str) -> str:
  prompt = '以“' + prompt + '”为主题，撰写一段' + style + '，字数在100字左右'
  print('功能：' + style + '，提示文案：' + prompt)
  
  response, history = model.chat(tokenizer, prompt, history=[])
    
  print('生成文案：' + response)

  return response


css_style = "#fixed_size_img {height: 240px;} "

title = "文案创作 by宁侠"
description = '''
本服务的主要应用场景涵盖多种文案输入生成和续写，例如用户可以自行输入各种内容，之后服务将会对其进行回答、续写或者按照指令进行回复。
'''

with gr.Blocks(title=title, css=css_style) as demo:
    gr.HTML('''
      <div style="text-align: center; max-width: 720px; margin: 0 auto;">
                  <div
                    style="
                      display: inline-flex;
                      align-items: center;
                      gap: 0.8rem;
                      font-size: 1.75rem;
                    "
                  >
                    <h1 style="font-family: PingFangSC; font-weight: 500; line-height: 1.5em; font-size: 32px; margin-bottom: 7px;">
                      文案创作
                    </h1>
                    <h1 style="font-family: PingFangSC; font-weight: 500; line-height: 1.5em; font-size: 16px; margin-bottom: 7px;">
                      by宁侠
                    </h1>
                  </div>
                </div>
      ''')

    gr.Markdown(description)
    with gr.Row():
        radio_style = gr.Radio(label="功能选择", choices=["小红书笔记", "小红书标题", "公众号文案", "朋友圈微商文案", "商品卖点", "商品描述", "商品种草文案", "商品好评", "广告标题", "创意广告", "产品起名", "视频拍摄剧本", "短视频口播稿", "直播脚本", "短视频拍摄提纲", "SEO文章", "产品slogan", "夸夸机器人", "诗词创作", "作文创作"], value="小红书笔记")
    with gr.Row():
        text_input = gr.Textbox(label="提示文案", value="探索西夏：沙漠风情与多元文化的西北之旅")
        text_output = gr.Textbox(label="AI创作")
    with gr.Row():
        btn_submit = gr.Button(value="一键创作", elem_id="blue_btn")
        # btn_clear = gr.Button(value="清除")

#    examples = gr.Examples(["以“个性iPhone手机壳”为主题，撰写一段朋友圈微商文案，字数在100字左右", "美化这句话：本服务主要应用于多种场景文案输入的生成和续写，比如用户可以自行尝试输入各种内容，然后让服务去回答、续写或者根据指令回复。"], inputs=[text_input], outputs=text_output)
    examples = gr.Examples(["探索西夏：沙漠风情与多元文化的西北之旅", "个性iPhone手机壳"], inputs=[text_input], outputs=text_output)
    btn_submit.click(text_generation, inputs=[text_input, radio_style], outputs=text_output)
    # btn_clear清除画布

demo.queue(api_open=False).launch(debug=True)