import gradio as gr
from transformers import pipeline
import pandas as pd
import json

pipe = pipeline("summarization", model="Gabriel/bart-base-cnn-xsum-swe")

def generate(in_text):
  print(in_text)
  answer = pipe(in_text, num_beams=5 ,min_length=20, max_length=120)
  print(answer)
  return answer[0]["summary_text"]
  
def update_history(df, in_text, gen_text ,generation_type, parameters):
    # get rid of first seed phrase
    new_row = [{"In_text": in_text,
                "Gen_text": gen_text,
                "Generation Type": generation_type,
                "Parameters": json.dumps(parameters)}]
    return pd.concat([df, pd.DataFrame(new_row)]) 
    
def generate_transformer(in_text, num_beams ,history):
    gen_text= generate(in_text)
    return gen_text, update_history(history, in_text, gen_text, "Transformer", {"num_beams": num_beams})
 

with gr.Blocks() as demo:
    gr.Markdown("""# Summarization Engine!""")

    with gr.Accordion("See Details", open=False):
        gr.Markdown("lorem ipsum")


    with gr.Tabs():
        with gr.TabItem("Transformer Generation"):
            gr.Markdown(
                      """The default parameters for distilgpt2 work well to generate moves. Use this tab as 
                      a baseline for your experiments.""")

            with gr.Row():
                with gr.Column(scale=4):
                    text_baseline_transformer= gr.Textbox(lines=4,label="Input Text", placeholder="hej hej",)
                with gr.Column(scale=3):
                    with gr.Row():
                        num_beams = gr.Slider(minimum=2, maximum=10, value=2, step=1, label="Number of Beams2")
                    output_basline_transformer = gr.Textbox(label="Output Text")
            transformer_button = gr.Button("Summarize!")            

        # with gr.TabItem("Strong Baseline"):
        #     gr.Markdown(
        #               """The default parameters for distilgpt2 work well to generate moves. Use this tab as 
        #               a baseline for your experiments.""")
        #     with gr.Row():
        #         with gr.Column(scale=4):
        #             text_baseline= gr.Textbox(lines=4,label="Input Text", placeholder="hej hej",)
        #         with gr.Column(scale=3):
        #             with gr.Row():
        #                 num_beams2 = gr.Slider(minimum=2, maximum=10, value=2, step=1, label="Number of Beams2")
        #                 num_beams3 = gr.Slider(minimum=2, maximum=10, value=2, step=1, label="Number of Beams3")
        #             output_basline = gr.Textbox(label="Output Text")
        #     baseline_button = gr.Button("Summarize!")

        # with gr.TabItem("LexRank"):
        #     gr.Markdown(
        #               """The default parameters for distilgpt2 work well to generate moves. Use this tab as 
        #               a baseline for your experiments.""")
        #     with gr.Row():
        #               label="Number of Beams")
        #         text_baseline= gr.Textbox(label="Input Text", placeholder="hej hej",)
        #         output_basline = gr.Textbox(label="Output Text")
        #     baseline_button = gr.Button("Summarize!")

    gr.Examples([["hi", 5]], [text_baseline_transformer, num_beams])       
    
    with gr.Box():
        gr.Markdown("<h3> Generation History <h3>")
        # Displays a dataframe with the history of moves generated, with parameters
        history = gr.Dataframe(headers=["In_text", "Gen_text", "Generation Type", "Parameters"], overflow_row_behaviour="show_ends", wrap=True)



    with gr.Box():
      	gr.Markdown("<h3>How did you make this?<h3>")
        # gr.Markdown("""hej bottom.""")
    transformer_button.click(generate_transformer, inputs=[text_baseline_transformer, num_beams ,history], outputs=[output_basline_transformer , history] )



   # baseline_button.click(generate_transformer, inputs=[text_baseline, num_beams2 ,history], outputs=[output_basline,history] )


demo.launch()