import gradio as gr
from transformers import pipeline

# pipeline_en = pipeline(task="text2text-generation", model="beyond/genius-large")
# pipeline_en = pipeline_zh 
pipeline_zh = pipeline(task="text2text-generation", model="beyond/genius-base-chinese")


def predict_en(sketch):
  # generated_text = pipeline_en(sketch, num_beams=3, do_sample=True, max_length=200)[0]['generated_text']
  # return generated_text
    return "The English model (`genius-large`) to too large to be maintained in a free space, please download the model checkpoint and run locally."

def predict_zh(sketch):
  generated_text = pipeline_zh(sketch, num_beams=3, do_sample=True, max_length=200)[0]['generated_text']
  return generated_text.replace(' ','')
  
 


with gr.Blocks() as demo:
    gr.Markdown("""
                # 💡GENIUS – generating text using sketches! 
                [Please check our GitHub repo for more details.](https://github.com/beyondguo/genius)
                
                We provide both English and Chinese GENIUS models.
                - For English version, the mask token is `<mask>`;
                - For Chinese version, the mask token is `[MASK]`.
                """)
    with gr.Tab("English"):
        input1 = gr.Textbox(lines=5, value="<mask> Conference on Empirical Methods <mask> submission of research papers <mask> Deep Learning <mask>")
        output1 = gr.Textbox(lines=5)
        button1 = gr.Button("Generate")
    with gr.Tab("Chinese"):
        input2 = gr.Textbox(lines=5, value="自然语言处理[MASK]谷歌公司[MASK]通用人工智能[MASK]")
        output2 = gr.Textbox(lines=5)
        output2 = output2
        button2 = gr.Button("Generate")

    # with gr.Accordion("Open for More!"):
    #     gr.Markdown("Look at me...")

    button1.click(predict_en, inputs=input1, outputs=output1)
    button2.click(predict_zh, inputs=input2, outputs=output2)

demo.launch()