juliensimon's picture
juliensimon HF staff
Initial version
6d85d11
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()