DATA_DIR="TODO: replace with your directory"

# Shard images

## COCO
# python -m src.datasets.images.shard_images \
#     --create \
#     --input_images_dir ${DATA_DIR}/static/raw/coco/\
#     --output_images_dir ${DATA_DIR}/static/instruct/coco/ \
#     --images_per_shard 512


## Visual Genome
# python -m src.datasets.images.shard_images \
#     --create \
#     --input_images_dir ${DATA_DIR}/static/raw/visual-genome/ \
#     --output_images_dir ${DATA_DIR}/static/instruct/visual-genome/ \
#     --images_per_shard 512


## ALLaVA LAION4V
# python -m src.datasets.images.shard_images \
#     --create \
#     --input_images_dir ${DATA_DIR}/static/raw/allava-laion\
#     --output_images_dir ${DATA_DIR}/static/instruct/allava-laion \
#     --images_per_shard 512

## Q-Instruct
# python -m src.datasets.images.shard_images \
#     --create \
#     --input_images_dir ${DATA_DIR}/static/raw/q-instruct-images \
#     --output_images_dir ${DATA_DIR}/static/instruct/q-instruct-images \
#     --images_per_shard 512

## SynthDogPL

# python -m src.datasets.images.shard_images \
#     --create \
#     --input_images_dir ${DATA_DIR}/static/raw/synthdog/SynthDoG_pl_500k/train \
#     --output_images_dir ${DATA_DIR}/static/instruct/synthdog-pl-500k \
#     --images_per_shard 512

## SynthDogEN

# python -m src.datasets.images.shard_images \
#     --create \
#     --input_images_dir ${DATA_DIR}/static/raw/synthdog/SynthDoG_en_100k/train \
#     --output_images_dir ${DATA_DIR}/static/instruct/synthdog-en-100k \
#     --images_per_shard 512


# Filter and process data

## LLaVA Instruct
# python -m src.datasets.images.shard_images \
#     --assign \
#     --input_images_dir ${DATA_DIR}/static/instruct/coco \
#     --input_path ${DATA_DIR}/processed/instruct/llava-instruct-150k/complex_reasoning_77k_85-15_pl_en_mixture.jsonl \
#     --output_path ${DATA_DIR}/processed/instruct/llava-instruct-150k/complex_reasoning_77k_85-15_pl_en_mixture.jsonl \

# python -m src.datasets.images.shard_images \
#     --assign \
#     --input_images_dir ${DATA_DIR}/static/instruct/coco \
#     --input_path ${DATA_DIR}/processed/instruct/llava-instruct-150k/conversation_58k_85-15_pl_en_mixture.jsonl \
#     --output_path ${DATA_DIR}/processed/instruct/llava-instruct-150k/conversation_58k_85-15_pl_en_mixture.jsonl \

# python -m src.datasets.images.shard_images \
#     --assign \
#     --input_images_dir ${DATA_DIR}/static/instruct/coco \
#     --input_path ${DATA_DIR}/processed/instruct/llava-instruct-150k/detail_23k_85-15_pl_en_mixture.jsonl \
#     --output_path ${DATA_DIR}/processed/instruct/llava-instruct-150k/detail_23k_85-15_pl_en_mixture.jsonl \

## ALLaVA LAION4V Instruct

DATASET_PATH="interim/allava"
DATASET_FILE="ALLaVA-Instruct-LAION-4V-towerplus-72B.jsonl"
DATASET_NAME=$(basename ${DATASET_FILE} .jsonl)


python -m src.datasets.translate \
    --input_path ${DATA_DIR}/${DATASET_PATH}/${DATASET_FILE} \
    --output_path ${DATA_DIR}/${DATASET_PATH}/${DATASET_NAME}-filtered.jsonl \
    --skip_translate --filter \
    --from_human_comet_threshold 0.6 \
    --from_gpt_comet_threshold 0.5

python -m src.datasets.resample \
    --first_path ${DATA_DIR}/${DATASET_PATH}/${DATASET_NAME}-filtered.jsonl \
    --second_path ${DATA_DIR}/raw/allava/ALLaVA-Instruct-LAION-4V.json \
    --output_path ${DATA_DIR}/processed/instruct/${DATASET_NAME}_85-15_pl_en_mixture.jsonl \
    --second_ratio 0.15 \
    --seed 42

python -m src.datasets.images.shard_images \
    --assign \
    --remove_without_image \
    --input_images_dir ${DATA_DIR}/static/instruct/allava-laion \
    --input_path ${DATA_DIR}/processed/instruct/${DATASET_NAME}_85-15_pl_en_mixture.jsonl \
    --output_path ${DATA_DIR}/processed/instruct/${DATASET_NAME}_85-15_pl_en_mixture.jsonl \

## LVIS
DATASET_PATH="interim/lvis-instruct4v"
DATASET_FILE="lvis_instruct4v_220k-towerplus-72B.jsonl"
DATASET_NAME=$(basename ${DATASET_FILE} .jsonl)


python -m src.datasets.translate \
    --input_path ${DATA_DIR}/${DATASET_PATH}/${DATASET_FILE} \
    --output_path ${DATA_DIR}/${DATASET_PATH}/${DATASET_NAME}-filtered.jsonl \
    --skip_translate --filter \
    --from_human_comet_threshold 0.5 \
    --from_gpt_comet_threshold 0.6

python -m src.datasets.resample \
    --first_path ${DATA_DIR}/${DATASET_PATH}/${DATASET_NAME}-filtered.jsonl \
    --second_path ${DATA_DIR}/raw/lvis-instruct4v/lvis_instruct4v_220k.json \
    --output_path ${DATA_DIR}/processed/instruct/${DATASET_NAME}_85-15_pl_en_mixture.jsonl \
    --second_ratio 0.15 \
    --seed 42

python -m src.datasets.images.shard_images \
    --assign \
    --input_images_dir ${DATA_DIR}/static/instruct/coco \
    --input_path ${DATA_DIR}/processed/instruct/${DATASET_NAME}_85-15_pl_en_mixture.jsonl \
    --output_path ${DATA_DIR}/processed/instruct/${DATASET_NAME}_85-15_pl_en_mixture.jsonl \

## Q-Instruct

DATASET_PATH="interim/q-instruct"
DATASET_FILE="q-instruct-towerplus-72B.jsonl"
DATASET_NAME=$(basename ${DATASET_FILE} .jsonl)


python -m src.datasets.translate \
    --input_path ${DATA_DIR}/${DATASET_PATH}/${DATASET_FILE} \
    --output_path ${DATA_DIR}/${DATASET_PATH}/${DATASET_NAME}-filtered.jsonl \
    --skip_translate --filter \
    --from_human_comet_threshold 0.8 \
    --from_gpt_comet_threshold 0.6

python -m src.datasets.resample \
    --first_path ${DATA_DIR}/${DATASET_PATH}/${DATASET_NAME}-filtered.jsonl \
    --second_path ${DATA_DIR}/raw/q-instruct/cleaned_labels.json \
    --output_path ${DATA_DIR}/processed/instruct/${DATASET_NAME}_85-15_pl_en_mixture.jsonl \
    --second_ratio 0.15 \
    --seed 42

python -m src.datasets.images.shard_images \
    --assign \
    --remove_without_image \
    --input_images_dir ${DATA_DIR}/static/instruct/q-instruct-images \
    --input_path ${DATA_DIR}/processed/instruct/${DATASET_NAME}_85-15_pl_en_mixture.jsonl \
    --output_path ${DATA_DIR}/processed/instruct/${DATASET_NAME}_85-15_pl_en_mixture.jsonl \

## A-OKVQA

DATASET_PATH="interim/aokvqa"
DATASET_FILE="aokvqa_v1p0_train_conversations-towerplus-72B.jsonl"
DATASET_NAME=$(basename ${DATASET_FILE} .jsonl)

python -m src.datasets.translate \
    --input_path ${DATA_DIR}/${DATASET_PATH}/${DATASET_FILE} \
    --output_path ${DATA_DIR}/${DATASET_PATH}/${DATASET_NAME}-filtered.jsonl \
    --skip_translate --filter \
    --from_human_comet_threshold 0.6

python -m src.datasets.resample \
    --first_path ${DATA_DIR}/${DATASET_PATH}/${DATASET_NAME}-filtered.jsonl \
    --second_path ${DATA_DIR}/interim/aokvqa/aokvqa_v1p0_train_conversations.jsonl \
    --output_path ${DATA_DIR}/processed/instruct/${DATASET_NAME}_85-15_pl_en_mixture.jsonl \
    --second_ratio 0.15 \
    --seed 42

python -m src.datasets.images.shard_images \
    --assign \
    --remove_without_image \
    --input_images_dir ${DATA_DIR}/static/instruct/coco \
    --input_path ${DATA_DIR}/processed/instruct/${DATASET_NAME}_85-15_pl_en_mixture.jsonl \
    --output_path ${DATA_DIR}/processed/instruct/${DATASET_NAME}_85-15_pl_en_mixture.jsonl \

# TallyQA

DATASET_PATH="interim/tallyqa"
DATASET_FILE="tallyqa_conversations-towerplus-72B.jsonl"
DATASET_NAME=$(basename ${DATASET_FILE} .jsonl)

python -m src.datasets.translate \
    --input_path ${DATA_DIR}/${DATASET_PATH}/${DATASET_FILE} \
    --output_path ${DATA_DIR}/${DATASET_PATH}/${DATASET_NAME}-filtered.jsonl \
    --skip_translate --filter \
    --from_human_comet_threshold 0.6

python -m src.datasets.resample \
    --first_path ${DATA_DIR}/${DATASET_PATH}/${DATASET_NAME}-filtered.jsonl \
    --second_path ${DATA_DIR}/interim/tallyqa/tallyqa_conversations.jsonl \
    --output_path ${DATA_DIR}/processed/instruct/${DATASET_NAME}_85-15_pl_en_mixture.jsonl \
    --second_ratio 0.15 \
    --seed 42

python -m src.datasets.images.shard_images \
    --assign \
    --input_images_dir ${DATA_DIR}/static/instruct/coco \
    --input_path ${DATA_DIR}/processed/instruct/${DATASET_NAME}_85-15_pl_en_mixture.jsonl \
    --output_path ${DATA_DIR}/processed/instruct/${DATASET_NAME}_85-15_pl_en_mixture.jsonl \

python -m src.datasets.images.shard_images \
    --assign \
    --remove_without_image \
    --input_images_dir ${DATA_DIR}/static/instruct/visual-genome \
    --input_path ${DATA_DIR}/processed/instruct/${DATASET_NAME}_85-15_pl_en_mixture.jsonl \
    --output_path ${DATA_DIR}/processed/instruct/${DATASET_NAME}_85-15_pl_en_mixture.jsonl \

# SynthDogPL

python -m src.datasets.synthdog.process_instruct \
    --input_path ${DATA_DIR}/static/raw/synthdog/SynthDoG_pl_500k/train/metadata.jsonl \
    --output_path ${DATA_DIR}/processed/instruct/synthdog_pl_500k_conversations.jsonl \
    --lang pl --image_source synthdog-pl-500k \
    --random_state 42

python -m src.datasets.images.shard_images \
    --assign \
    --input_images_dir ${DATA_DIR}/static/instruct/synthdog-pl-500k \
    --input_path ${DATA_DIR}/processed/instruct/synthdog_pl_500k_conversations.jsonl \
    --output_path ${DATA_DIR}/processed/instruct/synthdog_pl_500k_conversations.jsonl \

# SynthDogEN

python -m src.datasets.synthdog.process_instruct \
    --input_path ${DATA_DIR}/static/raw/synthdog/SynthDoG_en_100k/train/metadata.jsonl \
    --output_path ${DATA_DIR}/processed/instruct/synthdog_en_100k_conversations.jsonl \
    --lang en --image_source synthdog-en-100k \
    --random_state 42

python -m src.datasets.images.shard_images \
    --assign \
    --input_images_dir ${DATA_DIR}/static/instruct/synthdog-en-100k \
    --input_path ${DATA_DIR}/processed/instruct/synthdog_en_100k_conversations.jsonl \
    --output_path ${DATA_DIR}/processed/instruct/synthdog_en_100k_conversations.jsonl \