Update app.py

app.py

@@ -1,21 +1,27 @@
+
 import streamlit as st
-from transformers import AutoTokenizer, AutoModelWithLMHead
+from streamlit.elements.altair import generate_chart
+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 = AutoModelWithLMHead.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", "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)
-
+artist = st.text_input("Enter the artist", "Eminem")
+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)