#!/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 BUCKET_LOCATION=US

# 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.py"
    JAX_PLATFORMS=cpu python 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.py"
    JAX_PLATFORMS=cpu python 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}
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}/0/items

# 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}

JAX_PLATFORMS=cpu python MaxText/generate_param_only_checkpoint.py \
MaxText/configs/base.yml \
base_output_directory=${BASE_OUTPUT_DIRECTORY} \
load_parameters_path=${SCANNED_CKPT_PATH} \
run_name=${RUN_NAME} \
model_name=${MODEL_NAME} \
force_unroll=true
echo "Written MaxText unscanned checkpoint to ${BASE_OUTPUT_DIRECTORY}/${RUN_NAME}/checkpoints"

# 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
