File size: 2,391 Bytes
be08283 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 | import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from pathlib import Path
# Load model
MODEL_PATH = "outputs/final_model"
if Path(MODEL_PATH).exists():
tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_PATH)
model.eval()
MODEL_LOADED = True
else:
MODEL_LOADED = False
model = None
tokenizer = None
def predict_genre(track_name):
"""Predict genre for a track name"""
if not MODEL_LOADED:
return "Model not found. Please train first."
if not track_name:
return "Please enter a track name"
# Tokenize
inputs = tokenizer(track_name, return_tensors='pt', padding=True, truncation=True, max_length=256)
# Predict
with torch.no_grad():
outputs = model(**inputs)
probs = torch.softmax(outputs.logits, dim=-1)
pred_id = torch.argmax(probs, dim=-1).item()
confidence = probs[0, pred_id].item()
# Get label
pred_label = model.config.id2label.get(pred_id, f"Class_{pred_id}")
return f"**Genre:** {pred_label}\n\n**Confidence:** {confidence:.2%}"
# Create Gradio interface
with gr.Blocks(title="Spotify Genre Classifier", theme=gr.themes.Soft()) as demo:
gr.Markdown("# 🎵 Spotify Genre Classifier")
gr.Markdown("Enter a song track name to predict its genre using a fine-tuned GPT-2 model.")
with gr.Row():
with gr.Column():
track_input = gr.Textbox(
label="Track Name",
placeholder="e.g., Bohemian Rhapsody",
lines=1
)
predict_btn = gr.Button("🔮 Predict Genre", variant="primary")
with gr.Column():
output = gr.Textbox(label="Prediction")
# Examples
gr.Examples(
examples=[
"Bohemian Rhapsody",
"Shape of You",
"Old Town Road",
"Blinding Lights",
"Bad Guy",
"Stairway to Heaven",
"Smells Like Teen Spirit",
"Billie Jean",
],
inputs=track_input
)
predict_btn.click(fn=predict_genre, inputs=track_input, outputs=output)
track_input.submit(fn=predict_genre, inputs=track_input, outputs=output)
if __name__ == "__main__":
demo.launch()
|