import transformers
from transformers import pipeline
import gradio as gr

# Available models for pipeline
# checkpoint = 'wvangils/CTRL-Beatles-Lyrics-finetuned-newlyrics'
checkpoint = 'wvangils/GPT-Medium-Beatles-Lyrics-finetuned-newlyrics'
# checkpoint = 'wvangils/GPT-Neo-125m-Beatles-Lyrics-finetuned-newlyrics'
# checkpoint = 'wvangils/GPT2-Beatles-Lyrics-finetuned-newlyrics'
# checkpoint = 'wvangils/DistilGPT2-Beatles-Lyrics-finetuned-newlyrics'

# Create generator 
generator = pipeline("text-generation", model=checkpoint)

# Create function for generation
def generate_beatles(input_prompt, temperature):
    generated_lyrics = generator(input_prompt
                              , max_length = 100
                              , num_return_sequences = 1
                              , return_full_text = True
                              , verbose = 0
                              #, num_beams = 1
                              #, early_stopping = True    # Werkt niet goed lijkt 
                              , temperature = temperature    # Default 1.0    # Randomness, temperature = 1 minst risicovol, 0 meest risicovol
                              #, top_k = 50     # Default 50
                              , top_p = 0.5    # Default 1.0
                              , no_repeat_ngram_size = 3   # Default = 0
                              , repetition_penalty = 1.0   # Default = 1.0
                              #, do_sample = True     # Default = False
                              )[0]["generated_text"]
    return generated_lyrics

# Create textboxes for input and output
input_box = gr.Textbox(label="Input prompt:", placeholder="Write the start of a song here", lines=2)
output_box = gr.Textbox(label="Lyrics by The Beatles and GPT:", lines=20)
examples = [['In my dream I am', 0.7], ['I don\'t feel alive', 0.7]]
title='Beatles lyrics generator based on GPT2'
description='A medium class GPT2 model was fine-tuned on lyrics from The Beatles to generate Beatles-like text. Give it a try!'
temperature = gr.Slider(minimum=0.0, maximum=1.0, step=0.1, label="Temperature (high = sensitive for low probability tokens)", value=0.7, show_label=True)

# Use generate Beatles function in demo-app Gradio
gr.Interface(fn=generate_beatles
             , inputs=[input_box, temperature]
             , outputs=output_box
             , examples=examples
             , title=title
             , description=description
             ).launch()