import torch

from stable_audio_tools.data.dataset import MemmapDataset, GeneralMemmapMapDataset
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,
    # )

    dataset = GeneralMemmapMapDataset(
        "/app/suno/data/chirp_v4/vae/data_vae_val.bin",
        "/app/suno/data/chirp_v4/vae/data_semantic_val.bin",
        metas_path="/app/suno/data/chirp_v4/vae/metas_val.jsonl",
        vae_n_tokens=250,
        vae_dim=128,
        semantic_n_tokens=250,
    )

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