Spaces:
Running
Running
| import streamlit as st | |
| from streamlit.elements.altair import generate_chart | |
| from transformers import AutoTokenizer, AutoModelForCausalLM | |
| from transformers import pipeline | |
| title = st.title("Loading 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) | |
| text_generation = pipeline("text-generation", model=model, tokenizer=tokenizer) | |
| title.title("Rap lyrics generator") | |
| 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) |