import transformers
import gradio as gr
from transformers import pipeline

demo = gr.Blocks()

gpt2_5x4 = pipeline('text-generation', model = "Lifan-Z/Chinese-Classic-Poem-Generator-style5x4-GPT2")
gpt2_7x4 = pipeline('text-generation', model = "Lifan-Z/Chinese-Classic-Poem-Generator-style7x4-GPT2")
gpt2_5x8 = pipeline('text-generation', model = "Lifan-Z/Chinese-Classic-Poem-Generator-style5x8-GPT2")
gpt2_7x8 = pipeline('text-generation', model = "Lifan-Z/Chinese-Classic-Poem-Generator-style7x8-GPT2")


def poems_5x4(st):
  poems=""
  sequences = gpt2_5x4('<|endoftext|>'+st, max_length=26, do_sample=True, top_k=20, top_p=0.9, repetition_penalty=1.2, num_return_sequences=20, eos_token_id=0)
  i=0
  for seq in sequences:
    if (len(seq.get('generated_text')) == 60):
      poems = poems + (str(seq.get('generated_text'))[13:60]) + "\n\n"
      i=i+1
      if(i==6):
        return(poems)

def poems_7x4(st):
  poems=""
  sequences = gpt2_7x4('<|endoftext|>'+st, max_length=34, do_sample=True, top_k=20, top_p=0.9, repetition_penalty=1.2, num_return_sequences=20, eos_token_id=0)
  i=0
  for seq in sequences:
    if (len(seq.get('generated_text')) == 76):
      poems = poems + (str(seq.get('generated_text'))[13:76]) + "\n\n"
      i=i+1
      if(i==6):
        return(poems)


def poems_5x8(st):
  poems=""
  sequences = gpt2_5x8('<|endoftext|>'+st, max_length=50, do_sample=True, top_k=20, top_p=0.9, repetition_penalty=1.2, num_return_sequences=30, eos_token_id=0)
  i=0
  for seq in sequences:
    if (len(seq.get('generated_text')) == 108):
      poems = poems + (str(seq.get('generated_text'))[13:108]) + "\n\n"
      i=i+1
      if(i==6):
        return(poems)


def poems_7x8(st):
  poems=""
  sequences = gpt2_7x8('<|endoftext|>'+st, max_length=66, do_sample=True, top_k=20, top_p=0.9, repetition_penalty=1.2, num_return_sequences=40, eos_token_id=0)
  i=0
  for seq in sequences:
    if (len(seq.get('generated_text')) == 140):
      poems = poems + (str(seq.get('generated_text'))[13:140]) + "\n\n"
      i=i+1
      if(i==6):
        return(poems)



with demo:
    gr.Markdown("## **古诗辅助生成器**")
    gr.Markdown("### **生活不止眼前的苟且，还有诗和远方**")
    gr.Markdown("使用说明：五言古诗的平仄、用韵相当自由，也不限长短，这里只提供四句、八句两种。若想生成五言绝句、律诗，需要自己按平仄、用韵作出选择，再慢慢优化，逐句完成.\
     比如，要用'雨'字开头写一首五言律诗：先输入一个'雨'字，点击生成六首诗。假设只对第二首的头两句诗比较满意，则复制这两句作为输入。第二次再生成六首诗，这次对第五首的前六句满意，\
     再复制这六句作为输入。多次重复即可得到满意的结果。标点符号跟普通文字一样处理，需要注意的是末尾不要是空格。七言诗同样处理. --- Lf-Z 工作室")
    with gr.Tabs():
        with gr.TabItem("五言四句"):
            with gr.Row():
                text_input_5x4 = gr.Textbox(label="输入诗的开头文字，以英文空格分隔，末尾不用空格。一次六首，大概7秒。", placeholder="白 日 依 山 尽", lines=1)
                text_output_5x4 = gr.Textbox()
            button_5x4 = gr.Button("提交")
        with gr.TabItem("五言八句"):
            with gr.Row():
                text_input_5x8 = gr.Textbox(label="输入诗的开头文字，以英文空格分隔，末尾不用空格。一次六首，大概15秒。", placeholder="白 日 依 山 尽", lines=1)
                text_output_5x8 = gr.Textbox()
            button_5x8 = gr.Button("提交")
        with gr.TabItem("七言四句"):
            with gr.Row():
                text_input_7x4 = gr.Textbox(label="输入诗的开头文字，以英文空格分隔，末尾不用空格。一次六首，大概10秒。", placeholder="山 明 水 净 夜 来 霜", lines=1)
                text_output_7x4 = gr.Textbox()
            button_7x4 = gr.Button("提交")
        with gr.TabItem("七言八句"):
            with gr.Row():
                text_input_7x8 = gr.Textbox(label="输入诗的开头文字，以英文空格分隔，末尾不用空格。一次六首，大概30秒。", placeholder="山 明 水 净 夜 来 霜", lines=1)
                text_output_7x8 = gr.Textbox()
            button_7x8 = gr.Button("提交")



    button_5x4.click(poems_5x4, inputs=text_input_5x4, outputs=text_output_5x4)
    button_5x8.click(poems_5x8, inputs=text_input_5x8, outputs=text_output_5x8)
    button_7x4.click(poems_7x4, inputs=text_input_7x4, outputs=text_output_7x4)
    button_7x8.click(poems_7x8, inputs=text_input_7x8, outputs=text_output_7x8)

demo.launch(share=True)