Spaces:
Running
Running
| from os import CLD_CONTINUED | |
| import streamlit as st | |
| 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(): | |
| text_file = open("rappers.txt") | |
| rappers = text_file.readlines() | |
| rappers = [name[:-1] for name in rappers] | |
| rappers.sort() | |
| return rappers | |
| title = st.title("Loading model") | |
| tokenizer, model = load_model() | |
| text_generation = pipeline("text-generation", model=model, tokenizer=tokenizer) | |
| title.title("Rap lyrics generator") | |
| #artist = st.text_input("Enter the artist", "Wu-Tang Clan") | |
| list_of_rappers = load_rappers() | |
| artist = st.selectbox("Choose your rapper", tuple(list_of_rappers), index = len(list_of_rappers)-1) | |
| song_name = st.text_input("Enter the desired song name", "Sadboys") | |
| if st.button("Generate lyrics", help="The lyrics generation can last up to 2 minutres"): | |
| st.title(f"{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] | |
| for count, line in enumerate(generated_song['generated_text'].split("\n")): | |
| if"<EOS>" in line: | |
| break | |
| if count == 0: | |
| st.markdown(f"**{line[line.find('['):]}**") | |
| continue | |
| if "<BOS>" in line: | |
| st.write(line[5:]) | |
| continue | |
| if line.startswith("["): | |
| st.markdown(f"**{line}**") | |
| continue | |
| st.write(line) |