added predictors

ArtistCoherencyModel.py

@@ -3,6 +3,7 @@ from huggingface_hub import PyTorchModelHubMixin
 
 import torch
 import torch.nn as nn
+import torch.nn.functional as F
 
 from typing import Union
 
@@ -40,7 +41,7 @@ class ArtistCoherencyModel(nn.Module, PyTorchModelHubMixin):
             return self.artist_model(**inputs).logits
 
     def predict_artist(self, song: str) -> tuple[str, float]:
-        logits = self.generate_artist_logits(song)[0]
+        logits = F.softmax(self.generate_artist_logits(song)[0], dim=0)
         predicted_class_id = logits.argmax().item()
         return self.artist_model.config.id2label[predicted_class_id], float(
             logits[predicted_class_id]
@@ -54,7 +55,7 @@ class ArtistCoherencyModel(nn.Module, PyTorchModelHubMixin):
             return self.coherency_model(**inputs).logits
 
     def predict_coherency(self, song: str) -> tuple[str, float]:
-        logits = self.generate_coherency_logits(song)[0]
+        logits = F.softmax(self.generate_coherency_logits(song)[0], dim=0)
         predicted_class_id = logits.argmax().item()
         return self.coherency_model.config.id2label[predicted_class_id], float(
             logits[predicted_class_id]