// tensorflow-rnnlm.h

// Copyright (C) 2017 Intellisist, Inc. (Author: Hainan Xu)

// See ../../COPYING for clarification regarding multiple authors
//
// 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
//
// THIS CODE IS PROVIDED *AS IS* BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, EITHER EXPRESS OR IMPLIED, INCLUDING WITHOUT LIMITATION ANY IMPLIED
// WARRANTIES OR CONDITIONS OF TITLE, FITNESS FOR A PARTICULAR PURPOSE,
// MERCHANTABLITY OR NON-INFRINGEMENT.
// See the Apache 2 License for the specific language governing permissions and
// limitations under the License.

#ifndef KALDI_TFRNNLM_TENSORFLOW_RNNLM_H_
#define KALDI_TFRNNLM_TENSORFLOW_RNNLM_H_

#include <string>
#include <vector>
#include <unordered_map>
#include "util/stl-utils.h"
#include "base/kaldi-common.h"
#include "fstext/deterministic-fst.h"
#include "util/common-utils.h"

// Following macros are defined in both OpenFst and Tensorflow headers. Here we
// undef them before including "tensorflow/core/public/session.h" to silence
// compiler warnings. Note that this is not a panacea. We should still pay
// attention to the order of includes in other places in the codebase to avoid
// using the wrong macro definitions. Any OpenFst header or any header including
// an OpenFst header should be included before tfrnnlm/tensorflow-rnnlm.h. Also,
// to avoid macro redefinitions, any Tensorflow header should be included after
// tfrnnlm/tensorflow-rnnlm.h.
#undef LOG
#undef VLOG
#undef CHECK
#undef CHECK_EQ
#undef CHECK_LT
#undef CHECK_GT
#undef CHECK_LE
#undef CHECK_GE
#undef CHECK_NE
#undef DCHECK
#undef DCHECK_EQ
#undef DCHECK_LT
#undef DCHECK_GT
#undef DCHECK_LE
#undef DCHECK_GE
#undef DCHECK_NE

#include "tensorflow/core/public/session.h"

using tensorflow::Session;
using tensorflow::Tensor;

namespace kaldi {
namespace tf_rnnlm {

struct KaldiTfRnnlmWrapperOpts {
  std::string unk_symbol;
  int32 num_threads;  // 0 means unlimited

  KaldiTfRnnlmWrapperOpts() : unk_symbol("<oos>"), num_threads(1) {}

  void Register(OptionsItf *opts) {
    opts->Register("unk-symbol", &unk_symbol, "Symbol for out-of-vocabulary "
                   "words in rnnlm.");
    opts->Register("num-threads", &num_threads, "Number of threads for TF computation; "
                   "0 means unlimited.");
  }
};

/**
This class wraps the TensorFlow based RNNLM, and provides a set of interfaces
to be used for class TfRnnlmDeterministicFst, implemented later in this file
*/
class KaldiTfRnnlmWrapper {
 public:

  /// constructor
  /// opts specify symbol for <unk> and num-threads for computation
  /// rnn_wordlist specifies a wordlist file with format
  ///        [int-word-id] [word]
  ///     the word <oos> must appear in this file
  //
  /// word_symbol_table_rxfilename points to a standard word-list file in OpenFST style
  /// unk_prob_file has the format
  ///        [word]  [prob or count]  (it auto-normalizes the probabilities)
  /// tf_model_path is the location of the TensorFlow model
  KaldiTfRnnlmWrapper(const KaldiTfRnnlmWrapperOpts &opts,
                      const std::string &rnn_wordlist,
                      const std::string &word_symbol_table_rxfilename,
                      const std::string &unk_prob_file,
                      const std::string &tf_model_path);
  ~KaldiTfRnnlmWrapper() {
    session_->Close();
  }

  int32 GetEos() const { return eos_; }

  /// get an all-zero Tensor of the size that matches the hidden state of the TF model
  const Tensor& GetInitialContext() const;

  /// get the 2nd-to-last layer of RNN when feeding input of
  /// (initial-context, sentence-boundary)
  /// "cell" is short for "(last)cell-output"; calling it "cell" here because in
  /// later functions we have function GetLogProb() where we need to pass in
  /// one "cell" as input and another as output; to avoid confusing we use a single
  /// word "cell" for that instead of things like cell_out_in and cell_out_out.
  const Tensor& GetInitialCell() const;

  /// compute p(word | wseq) and return the log of that
  /// the computation used the input cell,
  /// which is the 2nd-to-last layer of the RNNLM associated with history wseq;
  ///
  /// and we generate (context_out, new_cell) by passing (context_in, word)
  /// into the TensorFlow session that manages the RNNLM
  /// if the last 2 pointers are NULL we don't query them in TF session
  /// e.g. in the case of computing p(</s>|some history)
  BaseFloat GetLogProb(int32 word,      // word id in RNN wordlist
                       int32 fst_word,  // FST word label, only for computing OOS cost
                       const Tensor &context_in,
                       const Tensor &cell_in,
                       Tensor *context_out,
                       Tensor *cell_out);

  /// takes in a word-id for FST and return the word-id for RNNLM
  /// return the word-id for <oos> if not found
  int FstLabelToRnnLabel(int i) const;

 private:
  /// read the TensorFlow model and create the session for computation
  /// num-threads need to be specified in creating the session
  void ReadTfModel(const std::string &tf_model_path, int32 num_threads);

  /// do queries on the session to get the initial tensors (cell + context)
  void AcquireInitialTensors();

  /// since usually we have a smaller vocab in RNN than the whole vocab,
  /// we use this mapping during rescoring
  std::vector<int> fst_label_to_rnn_label_;
  std::vector<std::string> rnn_label_to_word_;
  std::vector<std::string> fst_label_to_word_;

  KaldiTfRnnlmWrapperOpts opts_;
  Tensor initial_context_;
  Tensor initial_cell_;

  // this corresponds to the FST symbol table
  int32 num_total_words;
  // this corresponds to the RNNLM symbol table
  int32 num_rnn_words;

  Session* session_;  // for TF computation; pointer owned here
  int32 eos_;
  int32 oos_;

  std::vector<float> unk_costs_;  // extra cost for OOS symbol in RNNLM

  KALDI_DISALLOW_COPY_AND_ASSIGN(KaldiTfRnnlmWrapper);
};

class TfRnnlmDeterministicFst:
         public fst::DeterministicOnDemandFst<fst::StdArc> {
 public:
  typedef fst::StdArc::Weight Weight;
  typedef fst::StdArc::StateId StateId;
  typedef fst::StdArc::Label Label;

  // Does not take ownership.
  TfRnnlmDeterministicFst(int32 max_ngram_order, KaldiTfRnnlmWrapper *rnnlm);
  ~TfRnnlmDeterministicFst();
  void Clear();

  // We cannot use "const" because the pure virtual function in the interface is
  // not const.
  virtual StateId Start() { return start_state_; }

  // We cannot use "const" because the pure virtual function in the interface is
  // not const.
  virtual Weight Final(StateId s);

  virtual bool GetArc(StateId s, Label ilabel, fst::StdArc* oarc);

 private:
  typedef unordered_map<std::vector<Label>,
                        StateId, VectorHasher<Label> > MapType;
  StateId start_state_;
  MapType wseq_to_state_;
  std::vector<std::vector<Label> > state_to_wseq_;

  KaldiTfRnnlmWrapper *rnnlm_;
  int32 max_ngram_order_;
  std::vector<Tensor*> state_to_context_;
  std::vector<Tensor*> state_to_cell_;
};

}  // namespace tf_rnnlm
}  // namespace kaldi

#endif  // KALDI_TFRNNLM_TENSORFLOW_RNNLM_H_