import streamlit as st
from transformers import AutoTokenizer, AutoModelWithLMHead
from transformers import pipeline
st.title("Rap Lyrics Generator")
st.image('./parental.png')
model_ckpt = "flax-community/gpt2-rap-lyric-generator"
tokenizer = AutoTokenizer.from_pretrained(model_ckpt,from_flax=True)
model = AutoModelWithLMHead.from_pretrained(model_ckpt,from_flax=True)
text_generation = pipeline("text-generation", model=model, tokenizer=tokenizer)
artist = st.text_input("Enter the artist", "Kanye West")
prefix_text = st.text_input("Enter the prefix text", "Let's party tonight")
if len(artist)>0 and len(prefix_text)>0:
  prefix_text = "[Verse 1: "+artist+" ]" + "\n" + prefix_text
elif len(artist)>0:
   prefix_text = "[Verse 1: "+artist+" ]"
generated_text= text_generation(prefix_text, max_length=500, do_sample=True)[0]
output = generated_text['generated_text']
list1 = output.split("\n")
for l in list1:
  st.write(l)