add required files for demo

ArtistCoherencyModel.py

@@ -0,0 +1,73 @@
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+from huggingface_hub import PyTorchModelHubMixin
+
+import torch
+import torch.nn as nn
+
+from typing import Union
+
+from FFNN import FFNN
+
+
+class ArtistCoherencyModel(nn.Module, PyTorchModelHubMixin):
+    def __init__(self, config: dict):
+        super().__init__()
+
+        coherency_model_repo_id = config["coherency_model_repo_id"]
+        artist_model_repo_id = config["artist_model_repo_id"]
+        ffnn_model_repo_id = config["ffnn_model_repo_id"]
+
+        self.coherency_model_tokenizer = AutoTokenizer.from_pretrained(
+            coherency_model_repo_id
+        )
+        self.artist_model_tokenizer = AutoTokenizer.from_pretrained(
+            artist_model_repo_id
+        )
+
+        self.coherency_model = AutoModelForSequenceClassification.from_pretrained(
+            coherency_model_repo_id
+        )
+        self.artist_model = AutoModelForSequenceClassification.from_pretrained(
+            artist_model_repo_id
+        )
+        self.ffnn = FFNN.from_pretrained(ffnn_model_repo_id)
+
+    def generate_artist_logits(self, song: str) -> torch.FloatTensor:
+        inputs = self.artist_model_tokenizer(
+            song, return_tensors="pt", max_length=512, truncation=True
+        )
+        with torch.no_grad():
+            return self.artist_model(**inputs).logits
+
+    def generate_coherency_logits(self, song: str) -> torch.FloatTensor:
+        inputs = self.coherency_model_tokenizer(
+            song, return_tensors="pt", max_length=512, truncation=True
+        )
+        with torch.no_grad():
+            return self.coherency_model(**inputs).logits
+
+    def generate_song_embedding(self, song: str) -> torch.FloatTensor:
+        with torch.no_grad():
+            artist_logits = self.generate_artist_logits(song)
+            coherency_logits = self.generate_coherency_logits(song)
+            return torch.hstack((artist_logits[0], coherency_logits[0]))
+
+    def forward(self, song_or_embedding: Union[str, torch.Tensor]):
+        if type(song_or_embedding) is str:
+            song_or_embedding = self.generate_song_embedding(song_or_embedding)
+
+        return self.ffnn(song_or_embedding)
+
+    def generate_artist_coherency_logits(
+        self, song_or_embedding: Union[str, torch.Tensor]
+    ) -> torch.FloatTensor:
+        with torch.no_grad():
+            return self.forward(song_or_embedding)
+
+    def predict(
+        self, song_or_embedding: Union[str, torch.Tensor], return_ids: bool = False
+    ) -> Union[list[str], torch.Tensor]:
+        if type(song_or_embedding) is str:
+            song_or_embedding = self.generate_song_embedding(song_or_embedding)
+
+        return self.ffnn.predict(song_or_embedding, return_ids=return_ids)

FFNN.py

@@ -0,0 +1,89 @@
+from huggingface_hub import PyTorchModelHubMixin
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from typing import Union
+
+
+class FFNN(nn.Module, PyTorchModelHubMixin):
+    def __init__(self, config: dict) -> None:
+        super().__init__()
+
+        self.input_layer = nn.Linear(config["embedding_dim"], config["hidden_dim"])
+
+        self.hidden_layers = nn.ModuleList()
+        for layer_index in range(1, config["num_layers"]):
+            self.hidden_layers.append(
+                nn.Linear(config["hidden_dim"], config["hidden_dim"])
+            )
+
+        self.output_layer = nn.Linear(config["hidden_dim"], config["output_dim"])
+
+        self.id2label = config["id2label"]
+        self.label2id = config["label2id"]
+
+    def forward(self, embeddings: torch.Tensor) -> torch.Tensor:
+        z = F.relu(self.input_layer(embeddings))
+        for hidden_layer in self.hidden_layers:
+            z = F.relu(hidden_layer(z))
+        output = self.output_layer(z)
+        return F.softmax(output, dim=0)
+
+    def convert_logits_to_top_ids(self, logits: torch.Tensor) -> list[int]:
+        if len(logits.shape) != 1 and len(logits.shape) != 2:
+            raise ValueError("logits must either be a 1 or 2 dimensional tensor")
+
+        if len(logits.shape) == 1:
+            logits = [logits]
+
+        return [logits_row.argmax().item() for logits_row in logits]
+
+    def convert_logits_to_labels(self, logits: torch.Tensor) -> list[str]:
+        if len(logits.shape) != 1 and len(logits.shape) != 2:
+            raise ValueError("logits must either be a 1 or 2 dimensional tensor")
+
+        if len(logits.shape) == 1:
+            logits = [logits]
+
+        labels = []
+        for logits_row in logits:
+            labels.append(self.id2label[str(logits_row.argmax().item())])
+
+        return labels
+
+    def predict(
+        self, embeddings: torch.Tensor, return_ids: bool = False
+    ) -> Union[list[str], list[int]]:
+        if len(embeddings.shape) != 1 and len(embeddings.shape) != 2:
+            raise ValueError("embeddings must either be a 1 or 2 dimensional tensor")
+
+        with torch.no_grad():
+            logits = self.forward(embeddings)
+
+            if return_ids:
+                return self.convert_logits_to_top_ids(logits)
+
+            return self.convert_logits_to_labels(logits)
+
+    def generate_labeled_logits(self, embeddings: torch.Tensor) -> dict[str, float]:
+        if len(embeddings.shape) != 1 and len(embeddings.shape) != 2:
+            raise ValueError("embeddings must either be a 1 or 2 dimensional tensor")
+
+        with torch.no_grad():
+            logits = self.forward(embeddings)
+
+            if len(logits.shape) == 1:
+                logits = [logits]
+
+            labeled_logits_list = []
+
+            for logits_row in logits:
+                labeled_logits = {}
+                for id, logit in enumerate(logits_row):
+                    labeled_logits[self.id2label[str(id)]] = logit
+
+                labeled_logits_list.append(labeled_logits)
+
+            return labeled_logits_list

app.py

@@ -1,4 +1,14 @@
+from ArtistCoherencyModel import ArtistCoherencyModel
 import streamlit as st
+import pandas as pd
 
-x = st.slider("Select a value")
-st.write(x, "squared is", x * x)
+artists_df = pd.read_csv("artists.csv")
+artist_names_list = list(artists_df["name"])
+
+
+artist_name_input = st.selectbox("Artist", artist_names_list)
+st.write(artist_name_input)
+
+ensemble_model = ArtistCoherencyModel.from_pretrained(
+    "tjl223/artist-coherency-ensemble"
+)

artists.csv

@@ -0,0 +1,21 @@
+name, id
+Taylor Swift, taylor-swift
+Morgan Wallen, morgan-wallen
+Megan Thee Stallion, megan-thee-stallion
+Drake, drake
+Nicki Minaj, nicki-minaj
+Zach Bryan, zach-bryan
+Grateful Dead, grateful-dead
+Luke Combs, luke-combs
+21 Savage, 21-savage
+SZA, sza
+Olivia Rodrigo, olivia-rodrigo
+Chris Stapleton, chris-stapleton
+The Smile, the-smile
+Doja Cat, doja-cat
+Jack Harlow, jack-harlow
+Noah Kahan, noah-kahan
+Travis Scott, travis-scott
+Jelly Roll, jelly-roll
+The Weeknd, the-weeknd
+Dua Lipa, dua-lipa

requirements.txt

@@ -1,4 +1,5 @@
 transformers==4.40.0
 huggingface_hub==0.22.2
 torch==2.2.2
-numpy==1.26.4
+numpy==1.26.4
+pandas==2.2.2