import torch

from stable_audio_tools.data.dataset import MemmapDataset
from stable_audio_tools.configs.dataset_configs.custom_metadata.semantic import (
    get_custom_metadata,
)


if __name__ == "__main__":

    dataset = MemmapDataset(
        "/app/suno/data/chirp_v4/vae/data_vae_val.bin",
        semantic_memmap_path="/app/suno/data/chirp_v4/vae/data_semantic_val.bin",
        codec_memmap_path="/app/suno/data/chirp_v4/vae/data_codec_val.bin",
        custom_metadata_fn=get_custom_metadata,
    )

    for idx, (vae_embed, info) in enumerate(dataset):
        for key, val in info.items():
            if isinstance(val, torch.Tensor):
                print(key, val.shape)