#!/bin/bash
# Copyright 2024 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


# This script will do the following:
# - Create GCS buckets to store model artifacts for the JetStream Maxtext Inference demo.
# - Convert the downloaded checkpoints to MaxText compatible checkpoints.
# - Convert the MaxText compatible checkpoints to unscanned checkpoints for inference.
# Device requirements:
# - Both checkpoints conversion only requires CPU (with JAX CPU mode).
set -ex

idx=$(date +%Y-%m-%d-%H-%M)
# Modify the `MODEL` and `MODEL_VARIATION` based on the model you use.
export MODEL=$1
export MODEL_VARIATION=$2
export MODEL_NAME=${MODEL}-${MODEL_VARIATION}

# After downloading checkpoints, copy them to GCS bucket at $CHKPT_BUCKET
# Please use separate GCS paths for uploading open source model weights ($CHKPT_BUCKET) and MaxText compatible weights ($MODEL_BUCKET).
# Point these variables to a GCS bucket that you created.
# An example of CHKPT_BUCKET could be: gs://${USER}-maxtext/chkpt/${MODEL}/${MODEL_VARIATION}
export CHKPT_BUCKET=$3
export MODEL_BUCKET=$4

# Point `BASE_OUTPUT_DIRECTORY` to a GCS bucket that you created, this bucket will store all the files generated by MaxText during a run, specifically the unscanned checkpoint.
export BASE_OUTPUT_DIRECTORY=$5

export HUGGING_FACE_CHECKPOINT=$6

export LORA_INPUT_ADAPTERS_PATH=$7

export BUCKET_LOCATION=US

if [[ -z "HUGGING_FACE_CHECKPOINT" ]]; then
    echo "HUGGING_FACE_CHECKPOINT is required."
    exit 1
fi

# Create three GCS buckets for the demo.
gcloud storage buckets create ${MODEL_BUCKET} --location=${BUCKET_LOCATION} || true
gcloud storage buckets create ${BASE_OUTPUT_DIRECTORY} --location=${BUCKET_LOCATION} || true

# Convert model checkpoints to MaxText compatible checkpoints.
if [ "$MODEL" == "gemma" ]; then
    CONVERT_CKPT_SCRIPT="convert_gemma_chkpt"
    JAX_PLATFORMS=cpu python3 -m MaxText.${CONVERT_CKPT_SCRIPT} \
    --base_model_path ${CHKPT_BUCKET} \
    --maxtext_model_path ${MODEL_BUCKET}/${MODEL}/${MODEL_VARIATION}/${idx} \
    --model_size ${MODEL_VARIATION}
else
    # We install torch CPU because the checkpoint conversion script MaxText/llama_or_mistral_ckpt.py does not need a TPU/GPU.
    pip install torch --index-url https://download.pytorch.org/whl/cpu
    # llama_or_mistral_ckpt.py requires local path, so we need to copy the checkpoint from CHKPT_BUCKET to local.
    tmp_ckpt_path="/tmp/"
    gcloud storage cp -r ${CHKPT_BUCKET} ${tmp_ckpt_path}

    path_parts=(${CHKPT_BUCKET//\// })
    directory_substring=${path_parts[-1]}
    CONVERT_CKPT_SCRIPT="llama_or_mistral_ckpt"

    if [[ ! -z "${LORA_INPUT_ADAPTERS_PATH}" ]]; then
	lora_local_path="/tmp/"

	if [[ "${LORA_INPUT_ADAPTERS_PATH}" =~ ^gs:// ]]; then
	    path_parts=(${LORA_INPUT_ADAPTERS_PATH//\// })
	    lora_dir_substring=${path_parts[-1]}

	    lora_local_path="${tmp_ckpt_path}${lora_dir_substring}"
	    if [[ ! -d ${lora_local_path} ]]; then
		mkdir ${lora_local_path}
	    fi
	    gcloud storage cp -r ${LORA_INPUT_ADAPTERS_PATH} ${tmp_ckpt_path}
	else
	    lora_local_path=${LORA_INPUT_ADAPTERS_PATH}
	fi

	JAX_PLATFORMS=cpu python3 -m MaxText.${CONVERT_CKPT_SCRIPT} \
	--base-model-path ${tmp_ckpt_path}${directory_substring} \
	--maxtext-model-path ${MODEL_BUCKET}/${MODEL}/${MODEL_VARIATION}/${idx} \
	--model-size ${MODEL_NAME} \
	--lora-input-adapters-path ${lora_local_path} \
	--huggingface-checkpoint ${HUGGING_FACE_CHECKPOINT}
    else
	JAX_PLATFORMS=cpu python3 -m MaxText.${CONVERT_CKPT_SCRIPT} \
	--base-model-path ${tmp_ckpt_path}${directory_substring} \
	--maxtext-model-path ${MODEL_BUCKET}/${MODEL}/${MODEL_VARIATION}/${idx} \
	--model-size ${MODEL_NAME} \
	--huggingface-checkpoint ${HUGGING_FACE_CHECKPOINT}
    fi
fi
echo "Written MaxText compatible checkpoint to ${MODEL_BUCKET}/${MODEL}/${MODEL_VARIATION}/${idx}"

# We define `SCANNED_CKPT_PATH` to refer to the checkpoint subdirectory.
export SCANNED_CKPT_PATH=${MODEL_BUCKET}/${MODEL}/${MODEL_VARIATION}/${idx}

# Convert MaxText compatible checkpoints to unscanned checkpoints.
# Note that the `SCANNED_CKPT_PATH` is in a `scanned` format which is great for training but for efficient decoding performance we want the checkpoint in an `unscanned` format.
export RUN_NAME=${MODEL_NAME}_unscanned_chkpt_${idx}

if [[ ! -z "${LORA_INPUT_ADAPTERS_PATH}" ]]; then
    JAX_PLATFORMS=cpu python3 -m MaxText.generate_param_only_checkpoint \
    MaxText/configs/base.yml \
    base_output_directory=${BASE_OUTPUT_DIRECTORY} \
    load_parameters_path=${SCANNED_CKPT_PATH}/base/0/items \
    lora_input_adapters_path=${SCANNED_CKPT_PATH}/loras \
    run_name=${RUN_NAME} \
    model_name=${MODEL_NAME} \
    force_unroll=true
    echo "Written MaxText unscanned checkpoint to ${BASE_OUTPUT_DIRECTORY}/${RUN_NAME}/checkpoints"
else
    JAX_PLATFORMS=cpu python3 -m MaxText.generate_param_only_checkpoint \
    MaxText/configs/base.yml \
    base_output_directory=${BASE_OUTPUT_DIRECTORY} \
    load_parameters_path=${SCANNED_CKPT_PATH}/0/items \
    run_name=${RUN_NAME} \
    model_name=${MODEL_NAME} \
    force_unroll=true
    echo "Written MaxText unscanned checkpoint to ${BASE_OUTPUT_DIRECTORY}/${RUN_NAME}/checkpoints"
fi



# We will use the unscanned checkpoints by passing `UNSCANNED_CKPT_PATH` into `LOAD_PARAMETERS_PATH` in the following sections.
export UNSCANNED_CKPT_PATH=${BASE_OUTPUT_DIRECTORY}/${RUN_NAME}/checkpoints/0/items
