app.py

import streamlit as st
from transformers import AutoTokenizer, AutoModelForCausalLM
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 = AutoModelForCausalLM.from_pretrained(model_ckpt,from_flax=True)
text_generation = pipeline("text-generation", model=model, tokenizer=tokenizer)
artist = st.text_input("Enter the artist", "Jay-Z")
song_name = st.text_input("Enter the desired song name", "Gas is going")
if st.button("Generate lyrics"):
    st.title(f"{artist}: {song_name}")
    prefix_text = f"<BOS>{song_name} [Verse 1:{artist}]"
    generated_song = text_generation(prefix_text, max_length=500, do_sample=True)[0]
    for count, line in enumerate(generated_song['generated_text'].split("\n")):
      if count == 0:
        st.write(line[line.find('['):])
        continue
      if"<EOS>" in line:
        break
      if "<BOS>" in line:
        st.write(line[5:])
        continue
      st.write(line)