testing lyric generation

ArtistCoherencyModel.py

@@ -64,6 +64,19 @@ class ArtistCoherencyModel(nn.Module, PyTorchModelHubMixin):
         with torch.no_grad():
             return self.forward(song_or_embedding)
 
+    def generate_artist_coherency_score(
+        self, artist_name: str, song_or_embedding: Union[str, torch.Tensor]
+    ) -> float:
+        coherent_index = self.ffnn.label2id[f"{artist_name}-coherent"]
+        incoherent_index = self.ffnn.label2id[f"{artist_name}-incoherent"]
+        logits = self.generate_artist_coherency_logits(song_or_embedding)
+        coherent_score = logits[coherent_index]
+        incoherent_score = logits[incoherent_index]
+        score = (coherent_score + incoherent_score) * (
+            coherent_score / incoherent_score
+        )
+        return float(score)
+
     def predict(
         self, song_or_embedding: Union[str, torch.Tensor], return_ids: bool = False
     ) -> Union[list[str], torch.Tensor]:

LyricGeneratorModel.py

@@ -0,0 +1,23 @@
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+
+class LyricGeneratorModel:
+    def __init__(self, repo_id: str):
+        self.model = AutoModelForCausalLM.from_pretrained(repo_id)
+
+        self.tokenizer = AutoTokenizer.from_pretrained(repo_id)
+        self.tokenizer.truncation_side = "right"
+        self.tokenizer.pad_token = self.tokenizer.eos_token
+        self.tokenizer.padding_side = "right"
+
+    def generate_lyrics(self, prompt: str, max_length: int):
+        input_ids = self.tokenizer(prompt, return_tensors="pt").input_ids
+        input_ids = input_ids.to("cuda")
+
+        output_tokens = self.model.generate(
+            input_ids, do_sample=True, max_length=max_length
+        )
+
+        output_text = self.tokenizer.batch_decode(output_tokens)[0]
+
+        return output_text

app.py

@@ -2,6 +2,8 @@ from ArtistCoherencyModel import ArtistCoherencyModel
 import streamlit as st
 import pandas as pd
 
+from LyricGeneratorModel import LyricGeneratorModel
+
 artists_df = pd.read_csv("artists.csv")
 artist_names_list = list(artists_df["name"])
 
@@ -11,11 +13,15 @@ song_title = st.text_input("Song Title")
 song_description = st.text_area("Song Description")
 submit_button = st.button("Submit")
 
-if submit_button:
-    st.write(
-        f"[Song Title] {song_title}\n[Song Artist] {artist_name_input}\n[Song Description] {song_description}"
-    )
-
-ensemble_model = ArtistCoherencyModel.from_pretrained(
+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}"
+    lyrics = lyric_generator_model.generate_lyrics(prompt, 1000)
+    st.write(lyrics)