fixed bug in lyric generator parsing

.gitignore

@@ -1 +1,2 @@
-__pycache__/
+__pycache__/
+test.txt

LyricGeneratorModel.py

@@ -26,4 +26,7 @@ class LyricGeneratorModel:
 
         output_text = self.tokenizer.batch_decode(output_tokens)[0]
 
-        return output_text.split(" ->: ")[1]
+        if "->:" in output_text:
+            return output_text.split("->:")[1].strip()
+        else:
+            return output_text

app.py

@@ -14,7 +14,7 @@ def get_artists():
 @st.cache_resource
 def get_evaluator_model():
     lyric_evaluator_model = None
-    with st.spinner("Loading Evaluation Model"):
+    with st.spinner("Loading Evaluation Model..."):
         lyric_evaluator_model = ArtistCoherencyModel.from_pretrained(
             "tjl223/artist-coherency-ensemble"
         )
@@ -25,7 +25,7 @@ def get_evaluator_model():
 @st.cache_resource
 def get_generator_model():
     lyric_generator_model = None
-    with st.spinner("Loading Generator Model"):
+    with st.spinner("Loading Generator Model..."):
         lyric_generator_model = LyricGeneratorModel(
             "tjl223/testllama2-qlora-lyric-generator-with-description"
         )
@@ -45,8 +45,10 @@ song_description = st.text_area("Song Description")
 if st.button("Submit"):
     prompt = f"[Song Title] {song_title}\n[Song Artist] {artist_name_for_generator}\n[Song Description] {song_description}"
     print(f"Prompt: {prompt}")
-    st.write(prompt)
-    lyrics = lyric_generator_model.generate_lyrics(prompt, 1000)
+    lyrics = ""
+    with st.spinner("Generating Lyrics..."):
+        lyrics = lyric_generator_model.generate_lyrics(prompt, 1000)
+    st.success("Finished Generating Lyrics")
     print(f"Lyrics: {lyrics}")
     for line in lyrics.split("\n"):
         if line.startswith("["):