DATA_DIR="TODO: replace with your directory"

# Extract base data from raw dataset.
# echo "Processing raw WIT data... (takes a while...)"
# python -m src.datasets.wit.process_raw \
#     --input_dir $DATA_DIR/raw/wit/train \
#     --categories_path $DATA_DIR/raw/wit/wit_v1.category_tree.jsonl \
#     --output_path $DATA_DIR/interim/wit/wit_v1.pl.jsonl

# Perform train-test split
echo "Splitting train and validation sets..."
python -m src.datasets.wit.make_train_test_split \
    --input_path $DATA_DIR/interim/wit/wit_v1.pl.jsonl \
    --categories_path $DATA_DIR/raw/wit/wit_v1.category_tree.jsonl \
    --output_train_path $DATA_DIR/interim/wit/wit_v1.train.pl.jsonl \
    --output_test_path $DATA_DIR/interim/wit/wit_v1.test.pl.jsonl \
    --test_size 0.05 \
    --seed 42

# Create conversations
echo "Creating conversations..."
python -m src.datasets.wit.make_conversations \
    --input_path $DATA_DIR/interim/wit/wit_v1.train.pl.jsonl \
    --instructions_path $DATA_DIR/interim/wit/instructions.json \
    --output_path $DATA_DIR/interim/wit/wit_v1.train_conversations.pl.jsonl \
    --seed 42

# Assign shards
echo "Assigning shards..."
python -m src.datasets.images.shard_images \
    --assign \
    --remove_without_image \
    --input_images_dir ${DATA_DIR}/static/instruct/wit \
    --input_path $DATA_DIR/interim/wit/wit_v1.train_conversations.pl.jsonl \
    --output_path $DATA_DIR/processed/instruct/wit_v1.train.pl.jsonl \

# Verify
echo "Verifying..."
python -m src.datasets.verify \
    --input_path $DATA_DIR/processed/instruct/wit_v1.train.pl.jsonl \
    --images_dir ${DATA_DIR}/static/instruct \
    --workers 128