#! /bin/bash
  #
  # This script demonstrates a lexicon learning recipe, which aims to imrove
  # the pronounciation of abbreviated words in the TED-LIUM lexicon. It assumes
  # the model exp/tri3 already exists. Please see steps/dict/learn_lexicon_greedy.sh
  # for explanation of the options. 
  #
  # Copyright 2018  Xiaohui Zhang
  # Apache 2.0
  
  . ./cmd.sh
  . ./path.sh
  
  oov_symbol="<unk>"
  # The user may have an phonetisaurus-trained English g2p model ready.
  g2p_mdl_dir=
  # The dir which contains the reference lexicon (most probably hand-derived)
  # we want to expand/improve, and nonsilence_phones.txt,.etc which we need  
  # for building new dict dirs.
  ref_dict=data/local/dict
  # acoustic training data we use to get alternative
  # pronunciations and collet acoustic evidence.
  data=data/train
  # the cut-off parameter used to select pronunciation candidates from phone
  # decoding. We remove pronunciations with probabilities less than this value
  # after normalizing the probs s.t. the max-prob is 1.0 for each word."
  min_prob=0.1
  # Refer to steps/dict/select_prons_greedy.sh for the detailed meaning of
  # alpha, beta and delta. Basically, the three dimensions of alpha
  # and beta correspond to three pronunciation sources: phonetic-
  # decoding, G2P and the reference lexicon, and the larger a value is,
  # the more aggressive we'll prune pronunciations from that sooure.
  # The valid range of each dim. is [0, 1] (for alpha, and 0 means 
  # we never pruned pron from that source.) [0, 100] (for beta). 
  alpha="0.04,0.02,0"
  beta="30,5,0"
  # Floor value of the pronunciation posterior statistics.
  delta=0.00000001
  # This parameter determines how many pronunciations we keep for each word
  # after the first pass pruning. See steps/dict/internal/prune_pron_candidates.py
  # for details.
  vcr=16 
  
  # Intermediate outputs of the lexicon learning stage will be put into dir
  dir=exp/tri3_lex_greedy_work
  nj=35
  decode_nj=30
  stage=0
  lexlearn_stage=0
  affix="learned_greedy"
  
  . utils/parse_options.sh # accept options
  
  # The reference vocab is the list of words which we already have hand-derived pronunciations.
  ref_vocab=data/local/vocab.txt
  cat $ref_dict/lexicon.txt | awk '{print $1}' | sort | uniq > $ref_vocab || exit 1; 
  
  # Get a G2P generated lexicon for oov words (w.r.t the reference lexicon)
  # in acoustic training data.
  if [ $stage -le 0 ]; then
    if [ -z $g2p_mdl_dir ]; then
      g2p_mdl_dir=exp/g2p_phonetisaurus
      steps/dict/train_g2p_phonetisaurus.sh $ref_dict/lexicon.txt $g2p_mdl_dir || exit 1;
    fi
    awk '{for (n=2;n<=NF;n++) vocab[$n]=1;} END{for (w in vocab) printf "%s
  ",w;}' \
      $data/text | sort -u > $data/train_vocab.txt || exit 1;
    awk 'NR==FNR{a[$1] = 1; next} {if(!($1 in a)) print $1}' $ref_vocab \
      $data/train_vocab.txt | sort > $data/oov_train.txt || exit 1;
    steps/dict/apply_g2p_phonetisaurus.sh --nbest 5 $data/train_vocab.txt $g2p_mdl_dir \
      exp/g2p_phonetisaurus/lex_train || exit 1;
  fi
  
  # Learn a lexicon based on the acoustic training data and the reference lexicon.
  if [ $stage -le 1 ]; then
    steps/dict/learn_lexicon_greedy.sh --lexiconp-g2p "exp/g2p_phonetisaurus/lex_train/lexicon.lex" \
      --alpha $alpha --beta $beta --delta $delta \
      --min-prob $min_prob --cmd "$train_cmd" \
      --variant-counts-ratio $vcr \
      --stage $lexlearn_stage --nj 60 --oov-symbol $oov_symbol --retrain-src-mdl false \
      $ref_dict $ref_vocab $data exp/tri3 data/lang data/local/dict_${affix}_nosp \
      $dir || exit 1;
  fi
  
  # Add pronounciation probs to the learned lexicon.
  if [ $stage -le 2 ]; then
    utils/prepare_lang.sh --phone-symbol-table data/lang/phones.txt \
      data/local/dict_${affix}_nosp $oov_symbol data/local/lang_${affix}_nosp data/lang_${affix}_nosp || exit 1;
    
    steps/align_si.sh --nj $nj --cmd "$train_cmd" \
      $data data/lang_${affix}_nosp exp/tri2 exp/tri2_ali_${affix}_nosp || exit 1;
    
    steps/get_prons.sh --cmd "$train_cmd" data/train data/lang_${affix}_nosp exp/tri2_ali_${affix}_nosp || exit 1;
    
    utils/dict_dir_add_pronprobs.sh --max-normalize true \
      data/local/dict_${affix}_nosp exp/tri2_ali_${affix}_nosp/pron_counts_nowb.txt \
      exp/tri2_ali_${affix}_nosp/sil_counts_nowb.txt \
      exp/tri2_ali_${affix}_nosp/pron_bigram_counts_nowb.txt data/local/dict_${affix} || exit 1;
    
    utils/prepare_lang.sh --phone-symbol-table data/lang/phones.txt \
      data/local/dict_${affix} $oov_symbol data/local/lang_${affix} data/lang_${affix} || exit 1;
  fi
  
  # Re-decode
  if [ $stage -le 3 ]; then
    ! cmp data/lang_nosp/words.txt data/lang_${affix}/words.txt &&\
      echo "$0: The vocab of the affix lexicon and the reference vocab may be incompatible."
    cp data/lang_nosp/G.fst data/lang_${affix}/
    utils/mkgraph.sh data/lang_${affix} exp/tri3 exp/tri3/graph_${affix} || exit 1;
    
    for dset in dev test; do
    (  steps/decode_fmllr.sh --nj $decode_nj --cmd "$decode_cmd"  --num-threads 4 \
        exp/tri3/graph_${affix} data/${dset} exp/tri3/decode_${affix}_${dset} || exit 1;
    ) &
    done
  fi
  
  # RESULTS:
  # Baseline:
  # %WER 18.7 | 507 17783 | 83.9 11.4 4.7 2.6 18.7 92.3 | -0.006 | exp/tri3/decode_dev/score_17_0.0/ctm.filt.filt.sys
  # %WER 17.6 | 1155 27500 | 84.7 11.6 3.7 2.4 17.6 87.2 | 0.013 | exp/tri3/decode_test/score_15_0.0/ctm.filt.filt.sys
  
  # Re-decoding with the learned lexicon:
  # %WER 18.5 | 507 17783 | 84.3 11.2 4.5 2.8 18.5 92.3 | -0.007 | exp/tri3/decode_learned_greedy_dev/score_16_0.0/ctm.filt.filt.sys
  # %WER 17.5 | 1155 27500 | 84.9 11.5 3.6 2.4 17.5 87.5 | 0.035 | exp/tri3/decode_learned_greedy_test/score_14_0.0/ctm.filt.filt.sys
  
  # To see the effect to neural-net results, one should re-train NN with the learned lexicon.
  # Experiments have shown that, with the new lang dir, one should just re-run NN training
  # starting from the supervision generation (steps/align_fmllr_lats.sh) stage, and should
  # expect improved overall WERs and word recognition performance on words whose pronunciations
  # were changed.
  
  exit
  wait