import gradio
from transformers import pipeline
import torch

classifier = pipeline("text-classification", model="Veucci/lyric-to-genre")

def predict(inputs):

    result = classifier(inputs)[0]
    return_final = f"Predicted Genre: {str(result['label'])}\nScore: {str(result['score'])}"

    return return_final

example_lyrics = [
    "When I'm away from you, I'm happier than ever Wish I could explain it better I wish it wasn't true Give me a day or two to think of something clever To write myself a letter To tell me what to do, mm-mmm Do you read my interviews? Or do you skip my avenue? (My avenue) When you (when you) said you were passing through Was I even on your way? I knew when I asked you to (when I asked you to) Be cool about what I was telling you You'd do the opposite of what you said you'd do (what you said you'd do) And I'd end up more afraid Don't say it isn't fair You clearly weren't aware that you made me miserable So if you really wanna know",
    "Look, I was gonna go easy on you and not to hurt your feelings But I'm only going to get this one chance (six minutes, six minutes) Something's wrong, I can feel it (six minutes, six minutes, Slim Shady, you're on) Just a feeling I've got, like something's about to happen, but I don't know what If that means what I think it means, we're in trouble, big trouble And if he is as bananas as you say, I'm not taking any chances You are just what the doctor ordered I'm beginning to feel like a Rap God, Rap God All my people from the front to the back nod, back nod Now who thinks their arms are long enough to slap box, slap box? They said I rap like a robot, so call me Rapbot",
    "Come as you are, as you were As I want you to be As a friend, as a friend As an old enemy Take your time, hurry up Choice is yours, don't be late Take a rest as a friend As an old Memoria, memoria Memoria, memoria Come doused in mud, soaked in bleach As I want you to be As a trend, as a friend As an old Memoria, memoria Memoria, memoria"
]

gr_interface = gradio.Interface(
    fn=predict,
    allow_flagging='never',
    inputs=gradio.Textbox(
        placeholder="Enter a lyric",
        label="Lyric Input",
        style="width: 400px; height: 150px; font-size: 16px; padding: 10px;"
    ),
    outputs=gradio.Textbox(label="Predicted Genre", style="font-size: 18px;"),
    css="footer {visibility: hidden;}",
    title="Lyric Genre Classifier",
    description="The model was trained using the BERT language model on my song lyrics dataset to predict the genre of a given song based on its lyrics. This repository houses the machine learning model, which is capable of making predictions in three distinct genres: Pop, Rock, and Hip-Hop. For training and test codes check out Github page.",
    examples=example_lyrics,
    live=True,
    button_style="background-color: #4188c9; border-radius: 5px; color: white; font-size: 16px;"
)

gr_interface.launch()