import datasets
from renumics import spotlight
from renumics.spotlight import layout
from renumics.spotlight.layouts import debug_classification

dataset = datasets.load_dataset("renumics/speech_commands_enrichment_only", "raw_and_enrichment_combined")
stacked_ds = datasets.concatenate_datasets([dataset["train"], dataset["validation"], dataset["test"]])
print(dataset)
while True:
    layout = debug_classification(label='label_string', prediction='prediction', embedding='embedding_reduced', 
                                    features=["label", "prediction", "probability"], inspect={'audio': spotlight.Audio})
    dtypes = {
        "audio": spotlight.Audio,
        "embedding_reduced": spotlight.Embedding
    }
    view = spotlight.show(
        stacked_ds,
        dtype=dtypes,
        layout= layout
    )

    view.close()