export DATA_DIR="TODO: replace with your directory"


python -m src.datasets.wit.process_pretrain \
    --input_path ${DATA_DIR}/interim/wit/wit_v1.train_filtered.pl.jsonl \
    --output_path ${DATA_DIR}/processed/pretrain/wit_v1.train.pl.jsonl \
    --images_path ${DATA_DIR}/static/raw/wit-rescaled \
    --images_source wit \
    --random_state 42

rm -rf ${DATA_DIR}/static/pretrain/wit

python -m src.datasets.images.shard_images \
    --input_path ${DATA_DIR}/processed/pretrain/wit_v1.train.pl.jsonl \
    --input_images_dir ${DATA_DIR}/static/raw/wit-rescaled \
    --output_images_dir ${DATA_DIR}/static/pretrain/wit \
    --images_per_shard 512