#!/bin/bash
# shellcheck disable=SC2090,SC2086,SC2089,SC1091
# Default project path
PROJECT_PATH="$HOME/projects/mt_dao"

# Parse command-line options
if ! OPTIONS=$(getopt -o p: --long project_path: -n 'parse-options' -- "$@"); then
	echo "centralised_dataset_prep.sh: Error parsing options" >&2
	exit 1
fi

eval set -- "$OPTIONS"

while true; do
	case "$1" in
	-p | --project_path)
		PROJECT_PATH="$2"
		shift 2
		;;
	--)
		shift
		break
		;;
	*)
		break
		;;
	esac
done

echo "centralised_dataset_prep.sh: PROJECT_PATH=$PROJECT_PATH"
#! Moving to the project folder
cd /nfs-share/anon/projects/llm-foundry || exit
# GIT_LFS_SKIP_SMUDGE is necessary to avoid downloading large files
GIT_LFS_SKIP_SMUDGE=true uv sync

# DATA_ROOT=/home/anon/c4
# DATA_ROOT_MDS=/home/anon/mds-c4
DATA_ROOT_SMALL=/home/anon/my-copy-c4
# DATA_ROOT_SMALL_MDS=/home/anon/my-mds-copy-c4

#! Download and pre-processing
# #! Full dataset
# uv run python scripts/data_prep/convert_dataset_hf.py --dataset c4 --data_subset en --out_root $DATA_ROOT --splits train val --concat_tokens 2048 --tokenizer EleutherAI/gpt-neox-20b --eos_text '<|endoftext|>'
# uv run python scripts/data_prep/convert_dataset_hf.py --dataset c4 --data_subset en --out_root $DATA_ROOT_MDS --splits train val --concat_tokens 2048 --tokenizer EleutherAI/gpt-neox-20b --eos_text '<|endoftext|>' --compression zstd # Compressed
#! Small dataset
# uv run python scripts/data_prep/convert_dataset_hf.py --dataset c4 --data_subset en --out_root $DATA_ROOT_SMALL --splits train_small val_small --concat_tokens 2048 --tokenizer EleutherAI/gpt-neox-20b --eos_text '<|endoftext|>'
# uv run python scripts/data_prep/convert_dataset_hf.py --dataset c4 --data_subset en --out_root $DATA_ROOT_SMALL_MDS --splits train_small val_small --concat_tokens 2048 --tokenizer EleutherAI/gpt-neox-20b --eos_text '<|endoftext|>' --compression zstd # Compressed

#! Test data
# #! Full dataset
# uv run python llmfoundry/data/text_data.py --local_path $DATA_ROOT --split val
# uv run python llmfoundry/data/text_data.py --local_path $DATA_ROOT_MDS --split val # Compressed
#! Small dataset
uv run python llmfoundry/data/text_data.py --local_path "$DATA_ROOT_SMALL" --split val_small
# uv run python llmfoundry/data/text_data.py --local_path $DATA_ROOT_SMALL_MDS --split val_small # Compressed

#! Remove the positional arguments
eval set --
