text: #   Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
  #
  # 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.
  
  import os
  
  import paddle
  
  from model import CrossEntropyWithKL, VAESeq2SeqModel, Perplexity, NegativeLogLoss, TrainCallback
  from args import parse_args
  from data import create_data_loader
  
  
  def train(args):
      print(args)
      device = paddle.set_device(args.device)
      train_loader, dev_loader, test_loader, vocab, bos_id, pad_id, train_data_len = create_data_loader(
          args)
  
      net = VAESeq2SeqModel(embed_dim=args.embed_dim,
                            hidden_size=args.hidden_size,
                            latent_size=args.latent_size,
                            vocab_size=len(vocab) + 2,
                            num_layers=args.num_layers,
                            init_scale=args.init_scale,
                            enc_dropout=args.enc_dropout,
                            dec_dropout=args.dec_dropout)
  
      gloabl_norm_clip = paddle.nn.ClipGradByGlobalNorm(args.max_grad_norm)
  
      anneal_r = 1.0 / (args.warm_up * train_data_len / args.batch_size)
      cross_entropy = CrossEntropyWithKL(base_kl_weight=args.kl_start,
                                         anneal_r=anneal_r)
      model = paddle.Model(net)
  
      optimizer = paddle.optimizer.Adam(args.learning_rate,
                                        parameters=model.parameters(),
                                        grad_clip=gloabl_norm_clip)
  
      if args.init_from_ckpt:
          model.load(args.init_from_ckpt)
          print("Loaded checkpoint from %s" % args.init_from_ckpt)
  
      ppl_metric = Perplexity(loss=cross_entropy)
      nll_metric = NegativeLogLoss(loss=cross_entropy)
  
      model.prepare(optimizer=optimizer,
                    loss=cross_entropy,
                    metrics=[ppl_metric, nll_metric])
  
      model.fit(train_data=train_loader,
                eval_data=dev_loader,
                epochs=args.max_epoch,
                save_dir=args.model_path,
                shuffle=False,
                callbacks=[TrainCallback(ppl_metric, nll_metric, args.log_freq)],
                log_freq=args.log_freq)
  
      # Evaluation
      print('Start to evaluate on test dataset...')
      model.evaluate(test_loader, log_freq=len(test_loader))
  
  
  if __name__ == '__main__':
      args = parse_args()
      train(args)
  