import gradio as gr
from transformers import pipeline

def generate(text,the_model,max_length,temperature,num_beams,top_k,top_p,repetition_penalty):
    generator = pipeline('text-generation', model=the_model)
    result = generator(text, num_return_sequences=3,
    max_length=max_length,
    temperature=temperature,
    num_beams=num_beams,
    top_k=top_k,
    top_p=top_p,
    repetition_penalty = repetition_penalty,
    no_repeat_ngram_size=2,early_stopping=False)
    return result[0]["generated_text"],result[1]["generated_text"],result[2]["generated_text"]

demo = gr.Interface(
    fn=generate,
    inputs=[
        gr.Textbox(lines=5, label="Input Text"),
        gr.Dropdown(choices=['gpt2','gpt2-medium','gpt2-large','gpt2-xl'],value = 'gpt2',label="Choose model"),
        gr.Slider(value=50,label="Max Length",minimum=1,maximum=1000),
        gr.Slider(value=1.0,label="Temperature",minimum=0.0,maximum=1.0,step=0.05),
        gr.Slider(value=4,label="Num Beams",minimum=2,maximum=6,step=1),
        gr.Slider(value=90,label="Top-k",minimum=0,maximum=100),
        gr.Slider(value=0.9,label="Top-p",minimum=0.1,maximum=1,step=0.05),
        gr.Slider(value=1.1,label="Repetition penalty",minimum=0.2,maximum=2,step=0.1)

        ],
    outputs=[
        gr.Textbox(label="Generated Text 1"),
        gr.Textbox(label="Generated Text 2"),
        gr.Textbox(label="Generated Text 3")],
    title = "Text Generator GPT2 Pipeline",
    description = "Text Generator. \n Temperature control randomness, lowering results in less random completions. As approach the zero, the model becomes more repetitive."
)

demo.launch()