import torch
import numpy as np
import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification

repo_name = 'spotify/autonlp-huggingface-demo-song-lyrics-18923587'

tokenizer = AutoTokenizer.from_pretrained(repo_name)
model = AutoModelForSequenceClassification.from_pretrained(repo_name)
labels = model.config.id2label

def predict(lyrics):
	inputs = tokenizer(lyrics, padding=True, truncation=True, return_tensors="pt")
	outputs = model(**inputs)
	predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
	predictions = predictions.detach().numpy()[0]
	predictions = predictions*100
	sorted_indexes = np.argsort(predictions)
	return "These lyrics are {:.2f}% {}, {:.2f}% {} and {:.2f}% {}.".format(
		predictions[sorted_indexes[-1]], labels[sorted_indexes[-1]],
		predictions[sorted_indexes[-2]], labels[sorted_indexes[-2]],
		predictions[sorted_indexes[-3]], labels[sorted_indexes[-3]])

input = gr.inputs.Textbox(lines=20)
iface = gr.Interface(fn=predict, inputs=input, outputs="text")
iface.launch()