import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM
from transformers import pipeline

def load_model():
    model_ckpt = "flax-community/gpt2-rap-lyric-generator"
    tokenizer = AutoTokenizer.from_pretrained(model_ckpt, from_flax=True)
    model = AutoModelForCausalLM.from_pretrained(model_ckpt, from_flax=True)
    return tokenizer, model

def load_rappers():
    with open("rappers.txt") as text_file:
        rappers = text_file.readlines()
    rappers = [name.strip() for name in rappers]
    rappers.sort()
    return rappers

tokenizer, model = load_model()
text_generation = pipeline("text-generation", model=model, tokenizer=tokenizer)

def generate_lyrics(artist, song_name):
    prefix_text = f"<BOS>{song_name} [Verse 1:{artist}]"
    generated_song = text_generation(prefix_text, max_length=750, do_sample=True)[0]
    lyrics = ""
    for count, line in enumerate(generated_song['generated_text'].split("\n")):
        if "<EOS>" in line:
            break
        if count == 0:
            lyrics += f"**{line[line.find('['):]}**\n"
            continue
        if "<BOS>" in line:
            lyrics += f"{line[5:]}\n"
            continue
        if line.startswith("["):
            lyrics += f"**{line}**\n"
            continue
        lyrics += f"{line}\n"
    return lyrics

list_of_rappers = load_rappers()

interface = gr.Blocks(theme="Blane187/fuchsia")

with interface:
    gr.Markdown("# Rap Lyrics Generator")
    artist = gr.Dropdown(choices=list_of_rappers, label="Choose your rapper", value=list_of_rappers[-1])
    song_name = gr.Textbox(label="Enter the desired song name", value="Sadboys")
    generate_button = gr.Button("Generate lyrics")
    output = gr.Markdown()

    generate_button.click(generate_lyrics, inputs=[artist, song_name], outputs=output)

interface.launch()