adding evaluator

LyricGeneratorModel.py

@@ -1,4 +1,4 @@
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 
 from peft import PeftModel, PeftConfig
 
@@ -6,10 +6,11 @@ from peft import PeftModel, PeftConfig
 class LyricGeneratorModel:
     def __init__(self, repo_id: str):
         config = PeftConfig.from_pretrained(repo_id)
+        bnb_config = BitsAndBytesConfig(load_in_8bit=True)
         model = AutoModelForCausalLM.from_pretrained(
             config.base_model_name_or_path,
             return_dict=True,
-            load_in_8bit=True,
+            quantization_config=bnb_config,
             device_map="auto",
         )
         self.tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)

app.py

@@ -7,21 +7,32 @@ from LyricGeneratorModel import LyricGeneratorModel
 artists_df = pd.read_csv("artists.csv")
 artist_names_list = list(artists_df["name"])
 
+lyric_evaluator_model = None
+with st.spinner("Loading Evaluation Model"):
+    lyric_evaluator_model = ArtistCoherencyModel.from_pretrained(
+        "tjl223/artist-coherency-ensemble"
+    )
+st.success("Finished Loading Evaluation Model")
 
-artist_name_input = st.selectbox("Artist", artist_names_list)
+lyric_generator_model = None
+with st.spinner("Loading Generator Model"):
+    lyric_generator_model = LyricGeneratorModel(
+        "tjl223/testllama2-qlora-lyric-generator-with-description"
+    )
+st.success("Finished Loading Generator Model")
+
+artist_name = st.selectbox("Artist", artist_names_list)
 song_title = st.text_input("Song Title")
 song_description = st.text_area("Song Description")
 submit_button = st.button("Submit")
 
-lyric_evaluator_model = ArtistCoherencyModel.from_pretrained(
-    "tjl223/artist-coherency-ensemble"
-)
-
-lyric_generator_model = LyricGeneratorModel(
-    "tjl223/testllama2-qlora-lyric-generator-with-description"
-)
-
 if submit_button:
-    prompt = f"[Song Title] {song_title}\n[Song Artist] {artist_name_input}\n[Song Description] {song_description}"
+    prompt = f"[Song Title] {song_title}\n[Song Artist] {artist_name}\n[Song Description] {song_description}"
+    print(f"Prompt: {prompt}")
+    st.write(prompt)
     lyrics = lyric_generator_model.generate_lyrics(prompt, 1000)
+    print(f"Lyrics: {lyrics}")
     st.write(lyrics)
+    score = lyric_evaluator_model.generate_artist_coherency_score(artist_name, lyrics)
+    print(f"Score: {score}")
+    st.write(f"Score: {score}")