#!/usr/bin/env bash

MODELLG="$1"
DATALG="$2"
NAME="$3"

argslist=""
for (( i = 4; i <= $# ; i++ ))
  do
    j=${!i}
    argslist="${argslist} $j "
  done
echo $argslist >&2

cd "$(dirname $0)" || return

echo "Install fairseq" >&2
pip3 install -e fairseq -i https://bytedpypi.byted.org/simple/

prefix=hdfs://haruna/home/byte_arnold_lq_mlnlc/user/wangdanqing.122

dataset_path=${prefix}/Datasets/multilingual/data-bin/mspm/${DATALG}
tensorboard_logdir=${prefix}/Workshop/MultiLingual/${MODELLG}/logs/${NAME}
checkpoint_path=${prefix}/Workshop/MultiLingual/${MODELLG}/checkpoints/${NAME}
pretrained_path=${prefix}/Workshop/MultiLingual/pretrained

# change checkpoint path
hdfs dfs -mkdir -p $tensorboard_logdir
hdfs dfs -mkdir -p $checkpoint_path

local_root=~/transfer
resource_root=${local_root}/resource
output_path=${local_root}/output
model_path=${local_root}/model
mkdir -p ${resource_root}
mkdir -p ${output_path}
mkdir -p ${model_path}

local_dataset_path=${resource_root}/dataset
hadoop fs -copyToLocal ${dataset_path} ${local_dataset_path}
echo "Download resource from ${dataset_path} to ${local_dataset_path}" >&2

local_tensorboard_path=${output_path}/tensorboard_logdir
mkdir -p ${local_tensorboard_path}

local_checkpoint_path=${output_path}/checkpoint_path
mkdir -p ${local_checkpoint_path}
hadoop fs -copyToLocal ${checkpoint_path}/checkpoint_best.pt ${local_checkpoint_path}
echo "Load checkpoints from ${checkpoint_path}/checkpoint_best.pt to ${local_checkpoint_path}" >&2

local_pretrained_path=${model_path}/mbart.cc25
if [ ! -d ${local_pretrained_path} ]; then
  echo "Load pretrained model from ${pretrained_path}/mbart.CC25.tar.gz to ${local_pretrained_path}" >&2
  hadoop fs -copyToLocal ${pretrained_path}/mbart.CC25.tar.gz ${model_path}
  tar -xvzf ${model_path}/mbart.CC25.tar.gz -C ${model_path}
else
  echo "Pretrained model in ${local_pretrained_path}" >&2
fi

echo "Finish download files" >&2

langs=ar_AR,cs_CZ,de_DE,en_XX,es_XX,et_EE,fi_FI,fr_XX,gu_IN,hi_IN,it_IT,ja_XX,kk_KZ,ko_KR,lt_LT,lv_LV,my_MM,ne_NP,nl_XX,ro_RO,ru_RU,si_LK,tr_TR,vi_VN,zh_CN

echo "Transfer ${MODELLG} model to ${DATALG} test dataset..."
python fairseq/generate.py ${local_dataset_path}  \
--path ${local_checkpoint_path}/checkpoint_best.pt \
--task summarization_from_pretrained_mbart \
--gen-subset test \
--source-lang doc --target-lang sum \
--langs $langs \
--remove-bpe 'sentencepiece'  \
--min-len 50 \
--max-len-b 300 \
--lenpen 2 \
--no-repeat-ngram-size 3 \
--truncate-source \
$argslist \
> ${local_tensorboard_path}/${MODELLG}.${DATALG}.output

cat ${local_tensorboard_path}/${MODELLG}.${DATALG}.output | grep -P "^H" | sort -V |cut -f 3- > ${local_tensorboard_path}/${MODELLG}.${DATALG}.hypo

echo "Put ${local_tensorboard_path} to ${tensorboard_logdir}" >&2
hadoop fs -put -f ${local_tensorboard_path}/* ${tensorboard_logdir}/
sleep 600
