Spaces:
Running
Running
| from sklearn.preprocessing import StandardScaler | |
| from sklearn.model_selection import train_test_split | |
| import joblib | |
| import numpy as np | |
| import logging | |
| logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') | |
| logger = logging.getLogger(__name__) | |
| def dataset_splitter(X: np.ndarray, Y: np.ndarray): | |
| """ | |
| Script that splits the X and Y values to train, test, and valid splits. | |
| Parameters | |
| ---------- | |
| X : np.array | |
| Array of feature vectors | |
| Y : np.array | |
| Array of labels (real or fake) | |
| Returns | |
| ------- | |
| data : dict{np.array} | |
| A dictionary of np.arrays, containing the train/test/val split. | |
| """ | |
| logger.info(f"Dataset shape: {X.shape}, Labels: {len(Y)}") | |
| logger.info(f"Class distribution: {np.bincount(Y)}") | |
| # Split the data into train/val/test | |
| X_train, X_test, y_train, y_test = train_test_split( | |
| X, Y, test_size=0.1, random_state=42, stratify=Y | |
| ) | |
| X_train, X_val, y_train, y_val = train_test_split( | |
| X_train, y_train, test_size=0.2222, random_state=42, stratify=y_train | |
| ) | |
| logger.info(f"Train: {X_train.shape}, Validation: {X_val.shape}, Test: {X_test.shape}") | |
| data = { | |
| "train": (X_train, y_train), | |
| "val": (X_val, y_val), | |
| "test": (X_test, y_test), | |
| } | |
| return data | |
| def dataset_scaler(audio: np.ndarray, lyrics: np.ndarray): | |
| """ | |
| Method to scale both audio and lyric vectors using Z-Score. | |
| This allows us to have both vectors with a mean of 0, and ranges up and down based on the | |
| standard deviation without compromising the information it contains. | |
| This also saves the scalers through joblib, which will be loaded in the predict script. | |
| Parameters | |
| ---------- | |
| audio : np.array | |
| Array of audio features | |
| lyrics : np.array | |
| Array of lyric features | |
| Returns | |
| ------- | |
| scaled_audio : np.array | |
| Array of scaled audio features | |
| scaled_lyrics : np.array | |
| Array of scaled lyric features | |
| """ | |
| # Apply scalers to have similar-ranged data for both audio and lyrics training values | |
| audio_scaler = StandardScaler().fit(audio) | |
| lyric_scaler = StandardScaler().fit(lyrics) | |
| scaled_audio = audio_scaler.transform(audio) | |
| scaled_lyrics = lyric_scaler.transform(lyrics) | |
| # Save the trained scalers for prediction | |
| joblib.dump(audio_scaler, "models/fusion/audio_scaler.pkl") | |
| joblib.dump(lyric_scaler, "models/fusion/lyric_scaler.pkl") | |
| return scaled_audio, scaled_lyrics | |
| def instance_scaler(audio: np.ndarray, lyrics: np.ndarray): | |
| """ | |
| Method to scale the single input audio and lyrics | |
| Parameters | |
| ---------- | |
| audio : np.array | |
| Instance of an audio feature | |
| lyrics : np.array | |
| Instance of a lyric feature | |
| Returns | |
| ------- | |
| scaled_audio : np.array | |
| Array of scaled audio feature | |
| scaled_lyrics : np.array | |
| Array of scaled lyric feature | |
| """ | |
| # Apply scalers to the single inputs | |
| audio_scaler = joblib.load("models/fusion/audio_scaler.pkl") | |
| lyric_scaler = joblib.load("models/fusion/lyric_scaler.pkl") | |
| scaled_audio = audio_scaler.transform([audio]) | |
| scaled_lyrics = lyric_scaler.transform(lyrics) | |
| return scaled_audio, scaled_lyrics |