import numpy as np

data_dir = "/mnt/blob/CALVIN/task_ABCD_D/validation/"
lang_data = np.load(
        data_dir + "lang_annotations/auto_lang_ann.npy",
        allow_pickle=True,
    ).item()

ep_start_end_ids = lang_data["info"]["indx"]  # each of them are 64
lang_ann = lang_data["language"]["ann"]  # length total number of annotations
lang_task = lang_data["language"]["task"]

ann = {}
for i, (start_idx, end_idx) in enumerate(ep_start_end_ids):
    ann[lang_task[i]] = lang_ann[i]

import json

with open("examples/calvin/ann.json", "w") as f:
    json.dump(ann, f, indent=4)