import os
import gradio as gr
from transformers import pipeline

auth_token = os.environ.get("access_token")
pipeline_en = pipeline(task="text-classification", model="MeiJuice/CheckGPT",
                       use_auth_token=auth_token)
pipeline_zh = pipeline(task="text-classification", model="MeiJuice/CheckGPT-chinese",
                       use_auth_token=auth_token)


def predict_en(text):
    res = pipeline_en(text)[0]
    return res['label'], res['score']


def predict_zh(text):
    res = pipeline_zh(text)[0]
    return res['label'], res['score']


with gr.Blocks() as demo:
    with gr.Tab("English"):
        gr.Markdown("""
                    Note: Providing more text to the `Text` box can make the prediction more accurate!
                    """)
        t1 = gr.Textbox(lines=5, label='Text',
                        value="There are a few things that can help protect your credit card information from being misused when you give it to a restaurant or any other business:\n\nEncryption: Many businesses use encryption to protect your credit card information when it is being transmitted or stored. This means that the information is transformed into a code that is difficult for anyone to read without the right key.")
        button1 = gr.Button("🤖 Predict!")
        label1 = gr.Textbox(lines=1, label='Predicted Label 🎃')
        score1 = gr.Textbox(lines=1, label='Prob')
    with gr.Tab("中文版"):
        gr.Markdown("""
                    注意: 在`文本`栏中输入更多的文本，可以让预测更准确哦！
                    """)
        t2 = gr.Textbox(lines=5, label='文本',
                        value="对于OpenAI大力出奇迹的工作，自然每个人都有自己的看点。我自己最欣赏的地方是ChatGPT如何解决 “AI校正(Alignment)“这个问题。这个问题也是我们课题组这两年在探索的学术问题之一。")
        button2 = gr.Button("🤖 预测!")
        label2 = gr.Textbox(lines=1, label='预测结果 🎃')
        score2 = gr.Textbox(lines=1, label='模型概率')

    button1.click(predict_en, inputs=[t1], outputs=[label1, score1], api_name='predict_en')
    button2.click(predict_zh, inputs=[t2], outputs=[label2, score2], api_name='predict_zh')

    # # Page Count
    # gr.Markdown("""
    #             <center><a href='https://clustrmaps.com/site/1bsdc'  title='Visit tracker'><img src='//clustrmaps.com/map_v2.png?cl=080808&w=a&t=tt&d=NXQdnwxvIm27veMbB5F7oHNID09nhSvkBRZ_Aji9eIA&co=ffffff&ct=808080'/></a></center>
    #             """)

demo.launch()