import gradio as gr
from transformers import pipeline, GenerationConfig 

#Load pipeline
generator = pipeline("text-generation", model="gpt2")
config = GenerationConfig.from_pretrained("gpt2")

def greet(prompt, temperature, max_lenth, top_p, samples):
    #Enable greedy decoding if temperature is 0
    config.do_sample = True if temperature != 0 else False

    #Takes temperature and top_p value
    config.temperature = temperature
    config.top_p = top_p

    #Samples
    sample_input = {'Sample 1': 'Hi, this is a demo prompt for you',
                    'Sample 2': 'Alber Einstein is a famous physicist graduated from',
                    'Sample 3': 'University of Zurich located in'}
    
    #Choose between samples
    if samples and prompt == '':
        prompt = sample_input[samples]

    #Streaming the output
    for i in range(max_lenth):
        a = generator(prompt, max_new_tokens=1, generation_config=config)
        prompt=a[0]['generated_text']
        yield a[0]['generated_text']

demo = gr.Interface(
    fn=greet,
    inputs=[gr.Textbox(placeholder = "Write a tagline for an ice cream shop.", label="prompt", lines=5),
            gr.Slider(value=1, minimum=0, maximum=2, label='temperature',
                      info='''Temperature controls the randomness of the text generation.
                      1.0 makes the model more likely to generate diverse and sometimes more unexpected outputs.
                      0.0 makes the model's responses more deterministic and predictable.'''),
            gr.Slider(value=16, minimum=1, maximum=256, step=1, label='max_lenth',
                      info='''Maximum number of tokens that the model will generate in the output.'''),
            gr.Slider(value=1, minimum=0, maximum=1, label='top_p',
                      info='''Top-p controls the model's focus during text generation. 
                      It allows only the most probable tokens to be considered for generation, where the cumulative probability of these tokens must exceed this value.'''),
            gr.Dropdown(['Sample 1', 'Sample 2', 'Sample 3'], label="Sample Prompts",
                        info='''Some sample Prompts for you!''')],
    outputs=[gr.Textbox(label='Output texts')],
)

demo.launch()