test.py

# import webdataset as wds

# dataset_path  = "/home/mila/a/aditya.jain/scratch/global_model/webdataset/train/train450-000000.tar"

# # Create a WebDataset reader
# dataset = wds.WebDataset(dataset_path)

# for sample in dataset:
#     a = 2
#     for key, value in sample.items():
#         print(f"{key}: {type(value)}")

import json

categ_map_f = "/home/mila/a/aditya.jain/scratch/global_model/category_map.json"
new_categ_map = {}

with open(categ_map_f, "r") as f:
    category_map = json.load(f)

for key in category_map.keys():
    new_key = str(int(float(key)))
    new_categ_map[new_key] = category_map[key]


with open("/home/mila/a/aditya.jain/scratch/global_model/category_map_v2.json", "w") as f:
    json.dump(new_categ_map, f)