import gradio as gr
import random
import os
from transformers import AutoTokenizer, AutoModelForCausalLM

auth_token = os.environ['TOKEN'] or True

model = AutoModelForCausalLM.from_pretrained("aronvandepol/K-GPT125M", use_auth_token=auth_token)

tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neo-125M")

def get_gen(sample: str, length: int=30, beam: int =1):
    input_ids = tokenizer(sample, return_tensors='pt').input_ids
    output = model.generate(input_ids, 
        max_length = length, 
        num_beams = 5,
        no_repeat_ngram_size  = 3,
        early_stopping = True,
        do_sample=True,
        num_return_sequences = 10,
        pad_token_id=tokenizer.eos_token_id
    )
    generation = tokenizer.batch_decode(output, skip_special_tokens=True)
    gen_text = generation[beam-1]
    return gen_text


with gr.Blocks() as app:
    gr.Markdown("<h1><center>K-GPT_NEO</center></h1>")
    gr.Markdown(
        """
        <justify>
        Interact with the K-GPT_NEO model and generate kpop song texts! By entering a few words into the <i>input</i> prompt and press generate to get the most probable sentence. If you want to see some less probable results press the <i>I'm feeling lucky</i> button.
        </justify>
        """
        )
    beam=gr.Number(value=1, visible=False, precision=0)
    with gr.Row():
        length = gr.Slider(0, 100, step=5, label="Max generated words", value=30)
    with gr.Group():
        txt1 = gr.Textbox(label = "Input", placeholder="Type here and press enter...", lines=4)
        txt2 = gr.Textbox(label = "Output", placeholder="Generated sentence will appear here", lines=4, interactive=False)
    with gr.Row():
        btn = gr.Button("Generate most probable")
        rnd = gr.Button("Feeling Lucky!")
    btn.click(fn=get_gen, inputs=[txt1, length, beam], outputs=txt2)
    rnd.click(fn=get_gen, inputs=[txt1, length, gr.Number(value=random.randint(2, 10), visible=False, precision=0)], outputs=txt2)
    gr.Examples(examples=[['I miss you'], ['My Love has not faded yet'], ['Dancing the stars away']], inputs=txt1, outputs=txt2, fn=get_gen, cache_examples=True)
    gr.Markdown(
        """
        <br>
        <br>
        <br>
        <justify>
        K-GPT_NEO is based on the GPT-Neo text generation model developed by Eleuther AI. 
        This architecture was fine-tuned on 2000 English translations of K-pop songs ranging from BTS and BLACKPINK, to TWICE and ZONE. For more information on the training and data, please visit: 
        </justify>
        """
        )

app.launch()