#!/bin/bash
# shellcheck disable=SC2090,SC2086,SC2089,SC1091

# Default project path
PROJECT_PATH="$HOME/projects/mt_dao"
SAVE_PATH="$PROJECT_PATH/tokenizers_logs"
mkdir -p "$SAVE_PATH"

# Parse command-line options
if ! OPTIONS=$(getopt -o p: --long project_path: -n 'parse-options' -- "$@"); then
	echo "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 "PROJECT_PATH=$PROJECT_PATH"
#! Moving to the project folder
cd "$PROJECT_PATH" || exit
#! Preparing environment
if [[ $(hostname) == *'gpu-q'* ]]; then
	echo "Assuming the script is executing in the CSD3."
	#! Executing the environment preparation script
	#! NOTE: Must use "." to execute, "sh" doesn't work
	. "$PROJECT_PATH"/scripts/install_hpc_env.sh
else
	echo "Assuming the script is executing NOT in the CSD3."
	#! Executing the environment preparation script
	#! NOTE: Must use "." to execute, "sh" doesn't work
	. "$PROJECT_PATH"/scripts/install_env.sh
fi
#! Export the endpoint of the S3 object store, using directly the IP to avoid name resolution issues
export S3_ENDPOINT_URL='http://128.232.115.19:9000'
#! Exporting project folder so the script know where to store files
export PROJECT_PATH

# DONE
# DATETIME=$(date '+%Y%m%d_%H%M%S')
# python $PROJECT_PATH/mt_dao/dataset/train_tokenizer_from_tokenized.py \
# 	--output_folder_suffix "cosmo_$DATETIME" \
# 	--tokenizer "HuggingFaceTB/SmolLM-1.7B" \
# 	--decode_tokenizer "HuggingFaceTB/SmolLM-1.7B" \
# 	--splits train \
# 	--dataset_config_file "$PROJECT_PATH/mt_dao/conf/dataset/smollm-corpus-cosmo-shared.yaml" \
# 	--streams_config_file "$PROJECT_PATH/mt_dao/conf/dataset/streams/16_clients_single.yaml" \
# 	--batch_size 512 --num_workers 10 --truncate_num_batches 2000 \
# 	--vocab_size 49152 2>&1 | tee "$SAVE_PATH"/tokenizer_training_$DATETIME.log

# DONE
# DATETIME=$(date '+%Y%m%d_%H%M%S')
# python $PROJECT_PATH/mt_dao/dataset/train_tokenizer_from_tokenized.py \
# 	--output_folder_suffix "infiwebmath-4plus_$DATETIME" \
# 	--tokenizer "HuggingFaceTB/SmolLM-1.7B" \
# 	--decode_tokenizer "HuggingFaceTB/SmolLM-1.7B" \
# 	--splits train \
# 	--dataset_config_file "$PROJECT_PATH/mt_dao/conf/dataset/smollm-corpus-infiwebmath-4plus-shared.yaml" \
# 	--streams_config_file "$PROJECT_PATH/mt_dao/conf/dataset/streams/16_clients_single.yaml" \
# 	--batch_size 512 --num_workers 10 --truncate_num_batches 2000 \
# 	--vocab_size 49152 2>&1 | tee "$SAVE_PATH"/tokenizer_training_$DATETIME.log

# DONE
# DATETIME=$(date '+%Y%m%d_%H%M%S')
# python $PROJECT_PATH/mt_dao/dataset/train_tokenizer_from_tokenized.py \
# 	--output_folder_suffix "fine-math-4plus_$DATETIME" \
# 	--tokenizer "HuggingFaceTB/SmolLM-1.7B" \
# 	--decode_tokenizer "HuggingFaceTB/SmolLM-1.7B" \
# 	--splits train \
# 	--dataset_config_file "$PROJECT_PATH/mt_dao/conf/dataset/smollm-corpus-fine-math-4plus-shared.yaml" \
# 	--streams_config_file "$PROJECT_PATH/mt_dao/conf/dataset/streams/16_clients_single.yaml" \
# 	--batch_size 512 --num_workers 10 --truncate_num_batches 2000 \
# 	--vocab_size 49152 2>&1 | tee "$SAVE_PATH"/tokenizer_training_$DATETIME.log

# DONE
# DATETIME=$(date '+%Y%m%d_%H%M%S')
# python $PROJECT_PATH/mt_dao/dataset/train_tokenizer_from_tokenized.py \
# 	--output_folder_suffix "fine-web-edu-dedup_$DATETIME" \
# 	--tokenizer "HuggingFaceTB/SmolLM-1.7B" \
# 	--decode_tokenizer "HuggingFaceTB/SmolLM-1.7B" \
# 	--splits train \
# 	--dataset_config_file "$PROJECT_PATH/mt_dao/conf/dataset/smollm-corpus-fine-web-edu-dedup-shared.yaml" \
# 	--streams_config_file "$PROJECT_PATH/mt_dao/conf/dataset/streams/16_clients_single.yaml" \
# 	--batch_size 512 --num_workers 10 --truncate_num_batches 2000 \
# 	--vocab_size 49152 2>&1 | tee "$SAVE_PATH"/tokenizer_training_$DATETIME.log

# RUNNING on anon
DATETIME=$(date '+%Y%m%d_%H%M%S')
python $PROJECT_PATH/mt_dao/dataset/train_tokenizer_from_tokenized.py \
	--output_folder_suffix "math_datasets_$DATETIME" \
	--tokenizer "HuggingFaceTB/SmolLM-1.7B" \
	--decode_tokenizer "HuggingFaceTB/SmolLM-1.7B" \
	--splits train \
	--dataset_config_file "$PROJECT_PATH/mt_dao/conf/dataset/default.yaml" \
	--streams_config_file "$PROJECT_PATH/mt_dao/conf/dataset/streams/smollm_math_datasets_single.yaml" \
	--batch_size 512 --num_workers 192 --truncate_num_batches 1000 \
	--vocab_size 49152 2>&1 | tee "$SAVE_PATH"/tokenizer_training_$DATETIME.log

# RUNNING on anon
# DATETIME=$(date '+%Y%m%d_%H%M%S')
# python $PROJECT_PATH/mt_dao/dataset/train_tokenizer_from_tokenized.py \
# 	--output_folder_suffix "python-edu_$DATETIME" \
# 	--tokenizer "HuggingFaceTB/SmolLM-1.7B" \
# 	--decode_tokenizer "HuggingFaceTB/SmolLM-1.7B" \
# 	--splits train \
# 	--dataset_config_file "$PROJECT_PATH/mt_dao/conf/dataset/smollm-corpus-python-edu-shared.yaml" \
# 	--streams_config_file "$PROJECT_PATH/mt_dao/conf/dataset/streams/16_clients_single.yaml" \
# 	--batch_size 512 --num_workers 192 --truncate_num_batches 1000 \
# 	--vocab_size 49152 2>&1 | tee "$SAVE_PATH"/tokenizer_training_$DATETIME.log
