import gradio as gr
import numpy as np
import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer
from examples import yellow, stairway, numb, puppets, firework 

def lyrics_categories(input_text):
  spotify_model = "juliensimon/autonlp-song-lyrics-18753417"
  model = AutoModelForSequenceClassification.from_pretrained(spotify_model)
  tokenizer = AutoTokenizer.from_pretrained(spotify_model)
  labels = model.config.id2label
  inputs = tokenizer(input_text, return_tensors="pt")
  outputs = model(**inputs)
  predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
  predictions = predictions.detach().numpy()[0]
  index_sorted = np.argsort(predictions)[::-1]
  clean_outputs = {labels[idx]:str(predictions[idx]) for idx in index_sorted}
  print(clean_outputs)
  return clean_outputs
  
description = "With lyrics, find the top 5 genres this song belongs to! (Powered by Spotify)"
 
iface = gr.Interface(fn=lyrics_categories, 
                     inputs=gr.inputs.Textbox(lines=20, placeholder="Enter song lyrics here...", label="Song Lyrics"),
                     outputs=gr.outputs.Label(num_top_classes=5, label="Genres/Categories"),
                     examples=[stairway, numb, puppets, firework, yellow],
                     article=description,
                     title="Song Genre Predictor",
                    )
iface.launch()