added predictors

ArtistCoherencyModel.py

@@ -39,6 +39,13 @@ class ArtistCoherencyModel(nn.Module, PyTorchModelHubMixin):
         with torch.no_grad():
             return self.artist_model(**inputs).logits
 
+    def predict_artist(self, song: str) -> tuple[str, float]:
+        logits = self.generate_artist_logits(song)
+        predicted_class_id = logits.argmax().item()
+        return self.artist_model.config.id2label[predicted_class_id], float(
+            logits[predicted_class_id]
+        )
+
     def generate_coherency_logits(self, song: str) -> torch.FloatTensor:
         inputs = self.coherency_model_tokenizer(
             song, return_tensors="pt", max_length=512, truncation=True
@@ -46,6 +53,13 @@ class ArtistCoherencyModel(nn.Module, PyTorchModelHubMixin):
         with torch.no_grad():
             return self.coherency_model(**inputs).logits
 
+    def predict_coherency(self, song: str) -> tuple[str, float]:
+        logits = self.generate_artist_logits(song)
+        predicted_class_id = logits.argmax().item()
+        return self.coherency_model.config.id2label[predicted_class_id], float(
+            logits[predicted_class_id]
+        )
+
     def generate_song_embedding(self, song: str) -> torch.FloatTensor:
         with torch.no_grad():
             artist_logits = self.generate_artist_logits(song)

app.py

@@ -62,13 +62,26 @@ evaluator_title = st.title("Lyric Evaluator")
 artist_name_for_evaluator = st.selectbox(
     "Artist", artist_names_list, key="genorator_select"
 )
-song_lyrics = st.text_area("Song Lyrics")
+evaluator_song_lyrics = st.text_area("Song Lyrics")
 
-if st.button("Submit", key="genorator_submit"):
+if st.button("Submit", key="generator_submit"):
     score = lyric_evaluator_model.generate_artist_coherency_score(
-        artist_name_for_evaluator, song_lyrics.replace("\n\n", "\n")
+        artist_name_for_evaluator, evaluator_song_lyrics.replace("\n\n", "\n")
     )
-    print(song_lyrics)
     print(f"Score: {score}")
     st.write(f"Score: {score}")
-    print(song_lyrics.replace("\n\n", "\n"))
+
+predictor_title = st.title("Lyric Predictor")
+predictor_song_lyrics = st.text_area("Song Lyrics", key="predictor song lyrics")
+
+if st.button("Submit", key="predictor_submit"):
+    artist, artist_score = lyric_evaluator_model.predict_artist(
+        predictor_song_lyrics.replace("\n\n", "\n")
+    )
+    coherency, coherency_score = lyric_evaluator_model.predict_coherency(
+        predictor_song_lyrics.replace("\n\n", "\n")
+    )
+    print(f"Predicted {artist} with a score of {artist_score}")
+    st.write(f"Predicted {artist} with a score of {artist_score}")
+    print(f"Predicted {coherency} with a score of {coherency_score}")
+    st.write(f"Predicted {coherency} with a score of {coherency_score}")