inference-stat-sockeye.sh

#!/bin/bash

#
# Script for inference using Sockeye models, post-processing, and evaluation for quality and isometric translation.
# see required args below, update if exp setup is changed (model dir, test sets),
#	for isometric model inference pass arg '--isometric', along with nbest hypotheses re-ranking options.
#

# ---
set -e
DEVICE=0  # -1 for inference on cpu
export CUDA_VISIBLE_DEVICES=$DEVICE

SRC=en
TGT=$1
TEST_SET="mustc" 	# test set prefix [mustc, blind]

WRKDIR=../   # `./isometric-slt/`
MT_LC_STAT=$WRKDIR/scripts/compute_mt_isometric_stat.sh
JSON=true 					                        # inference input format
GEN_JSON=$WRKDIR/scripts/gen_json_file.py 	# script for generating json file
SP=$WRKDIR/scripts/sentence_piece.py 		    # script for post-processing

# inference arguments
INPUT_BPE=$2  	# inference input
REF=$3 		      # (optional) detoknized reference for evaluation
SPMODEL=$4	    # sentencepiece model for post-processing
MODELDIR=$5  		# sockeye models dir
CHECKPOINT='params_best.pt'		            # best checkpoint for inference
EVALDIR=$MODELDIR/$TEST_SET/evaluation		# inference out and log dir

# isometric inference arguments
ISOMETRIC=$6			      # "--isometric" - inference with isometric mt model
RERANK=$7			          # generate nbest list & rerank options=['isometric-lc', 'isometric-ratio', 'isometric-diff']
ISOMETRIC_ALPHA=0.7		      # alpha for reranking with *-ratio/diff - tuned per lang pair
ISOMETRIC_TOK="<2normal>" 	# model is trained with isometric tags = [2short, 2normal, 2long]

mkdir -p $EVALDIR && echo "DECODING IN: $EVALDIR"
MODEL=${MODELDIR}   #/$CHECKPOINT 	# Sockeye captures best checkpoint
echo "MODEL: $MODEL"
# ---


# input fo inference
cp $INPUT_BPE $EVALDIR/
INPUT_FILE="$(basename $INPUT_BPE)"
if $JSON; then
  INPUT=$EVALDIR/$INPUT_FILE.json
  if [ "$ISOMETRIC" = "--isometric" ]; then
    python $GEN_JSON -i $EVALDIR/$INPUT_FILE -o $INPUT -sp $ISOMETRIC_TOK
  else
    python $GEN_JSON -i $EVALDIR/$INPUT_FILE -o $INPUT 
  fi
else
  INPUT=$EVALDIR/$INPUT_FILE
fi
echo "INPUT: $INPUT"


# inference
if [ -z "$RERANK" ]; then
  LENPEN=1; BATCH=64; BEAM=5; NBEST=1
else
  LENPEN=1; BATCH=8; BEAM=5; NBEST=5
fi 
OUTPUT=$INPUT.lp$LENPEN.beam$BEAM.nbest$NBEST.$TGT
if [ ! -f "$OUTPUT" ]; then
	echo "OUTPUT: $OUTPUT"
	STARTTIME=$(date +%s)
	if $JSON; then
    sockeye-translate --models $MODEL \
      --json-input --input $INPUT \
      --output $OUTPUT \
      --batch-size $BATCH --beam-size $BEAM \
      --length-penalty-alpha $LENPEN  --nbest-size $NBEST \
	else
    sockeye-translate --models $MODEL \
        --input $INPUT_BPE \
        --output $OUTPUT \
        --batch-size $BATCH --beam-size $BEAM \
        --length-penalty-alpha $LENPEN --nbest-size $NBEST \
	fi
	ENDTIME=$(date +%s)
fi


# (optional) nbest hypotheses reranking for isometric translation
if [ ! -z "$RERANK" -a -f "$OUTPUT" -a "$JSON" = "true" ]; then
  echo "RERANKING: $OUTPUT"
  # re-rank hypotheses with isometric criterion
  OUTPUT_PRERANK=$OUTPUT.$RERANK && cp $OUTPUT $OUTPUT_PRERANK
  OUTPUT=$OUTPUT_PRERANK.rerank
  sockeye-rerank -hy $OUTPUT_PRERANK -r $REF -m $RERANK --isometric-alpha $ISOMETRIC_ALPHA \
	  --output-best --output-best-non-blank > $OUTPUT
fi


# pos-processing and evaluation 
LOG=$OUTPUT.stat
if [ -f "$OUTPUT" -a -f "$REF" -a ! -f "$LOG" ]; then
    OUTPUT_DETOK=$OUTPUT.detok
    python $SP --run "decode" --sp_model $SPMODEL --inputs $OUTPUT  --outputs $OUTPUT_DETOK

    echo "METRIC LOG: $LOG"
    $MT_LC_STAT $SRC $TGT $INPUT_BPE $OUTPUT_DETOK $REF $LOG
    echo -e "======\n" | tee -a $LOG
else
	echo "Skipping postprocessing & evaluation: $LOG"
fi