What Makes a Song
Go Viral?

# 1. Install: pip install kagglehub[pandas-datasets] pandas
import kagglehub
from kagglehub import KaggleDatasetAdapter
import pandas as pd

# Set the path to the file you'd like to load
file_path = "tracks.csv"

# Load the latest version
df = kagglehub.load_dataset(
    KaggleDatasetAdapter.PANDAS,
    "vedauv/what-makes-a-song-go-viral-10000-tracks",
    file_path,
)

print("Shape:", df.shape)
print(df.columns.tolist())
print("First 5 records:", df.head())

import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np

audio_features = ['danceability', 'energy', 'key', 'loudness',
                  'mode', 'speechiness', 'acousticness',
                  'instrumentalness', 'liveness', 'valence',
                  'tempo', 'duration_ms', 'time_signature']

df['log_streams'] = np.log1p(df['streams'])

plt.figure(figsize=(11, 9))
sns.heatmap(df[audio_features + ['log_streams']].corr(),
            annot=True, fmt='.2f', cmap='coolwarm', center=0)
plt.title('Audio Features vs Log(Streams) — Correlation')
plt.tight_layout()
plt.show()

import numpy as np
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split

# Drop rows missing critical columns
df = df.dropna(subset=audio_features + ['streams'])

# Log-transform streams (kills right-skew)
df['log_streams'] = np.log1p(df['streams'])

# Define "viral" as top 25% of streams
viral_threshold = df['streams'].quantile(0.75)
df['is_viral'] = (df['streams'] >= viral_threshold).astype(int)

print("Viral threshold:", int(viral_threshold), "streams")
print("Viral rate:", df['is_viral'].mean())

# Train/test split for both tasks
X = df[audio_features]
y_reg = df['log_streams']
y_clf = df['is_viral']

X_train, X_test, y_reg_train, y_reg_test = train_test_split(
    X, y_reg, test_size=0.2, random_state=42)
Xc_train, Xc_test, y_clf_train, y_clf_test = train_test_split(
    X, y_clf, test_size=0.2, random_state=42, stratify=y_clf)

# Scale features (helps even tree models for fair importance ranking)
scaler = StandardScaler()
X_train_s = scaler.fit_transform(X_train)
X_test_s = scaler.transform(X_test)

from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_absolute_error, r2_score
import numpy as np

reg = RandomForestRegressor(
    n_estimators=300,
    max_depth=None,
    min_samples_leaf=5,
    random_state=42,
    n_jobs=-1,
)

reg.fit(X_train_s, y_reg_train)
y_pred = reg.predict(X_test_s)

print("R²:", r2_score(y_reg_test, y_pred))
print("MAE (log scale):", mean_absolute_error(y_reg_test, y_pred))

# Convert MAE back to stream-count units
mae_streams = np.expm1(mean_absolute_error(y_reg_test, y_pred))
print("MAE ≈ %d streams" % mae_streams)

from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import (classification_report,
                                             confusion_matrix,
                                             roc_auc_score)
import seaborn as sns
import matplotlib.pyplot as plt

Xc_train_s = scaler.fit_transform(Xc_train)
Xc_test_s = scaler.transform(Xc_test)

clf = RandomForestClassifier(
    n_estimators=300,
    max_depth=None,
    min_samples_leaf=5,
    class_weight='balanced',
    random_state=42,
    n_jobs=-1,
)

clf.fit(Xc_train_s, y_clf_train)
yc_pred = clf.predict(Xc_test_s)
yc_proba = clf.predict_proba(Xc_test_s)[:, 1]

print(classification_report(y_clf_test, yc_pred,
                        target_names=['Not viral', 'Viral']))
print("ROC-AUC:", roc_auc_score(y_clf_test, yc_proba))

# Confusion Matrix
cm = confusion_matrix(y_clf_test, yc_pred)
plt.figure(figsize=(6, 5))
sns.heatmap(cm, annot=True, fmt='d',
            xticklabels=['Not viral', 'Viral'],
            yticklabels=['Not viral', 'Viral'],
            cmap='Blues')
plt.xlabel('Predicted')
plt.ylabel('Actual')
plt.title('Viral vs Not Viral — Confusion Matrix')
plt.tight_layout()
plt.show()

import pandas as pd
import matplotlib.pyplot as plt

# Get feature importances from the classifier
importance = pd.Series(
    clf.feature_importances_,
    index=audio_features
).sort_values(ascending=True)

print("What makes a song go viral? (ranked)")
print(importance.sort_values(ascending=False))

# Plot
plt.figure(figsize=(9, 6))
importance.plot(kind='barh', color='#22d3ee')
plt.title('What Makes a Song Go Viral? — Feature Importance')
plt.xlabel('Importance (higher = more predictive of virality)')
plt.tight_layout()
plt.show()

# Partial dependence: how does danceability affect viral probability?
from sklearn.inspection import partial_dependence
pd_result = partial_dependence(
    clf, Xc_test_s, features=['danceability'], kind='average'
)
print("Top 3 viral predictors:", importance.tail(3).index.tolist())