text: # coding=utf-8
  # Copyright (C) 2019 Alibaba Group Holding Limited
  #
  # 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 math
  import random
  import torch
  import torch.nn.functional as f
  from tqdm import tqdm
  from .network import Network
  from .utils.metrics import registry as metrics
  
  
  class Model:
      prefix = 'checkpoint'
      best_model_name = 'best.pt'
  
      def __init__(self, args, state_dict=None):
          self.args = args
  
          # network
          self.network = Network(args)
          self.device = torch.cuda.current_device() if args.cuda else torch.device('cpu')
          self.network.to(self.device)
          # optimizer
          self.params = list(filter(lambda x: x.requires_grad, self.network.parameters()))
          self.opt = torch.optim.Adam(self.params, args.lr, betas=(args.beta1, args.beta2),
                                      weight_decay=args.weight_decay)
          # updates
          self.updates = state_dict['updates'] if state_dict else 0
  
          if state_dict:
              new_state = set(self.network.state_dict().keys())
              for k in list(state_dict['model'].keys()):
                  if k not in new_state:
                      del state_dict['model'][k]
              self.network.load_state_dict(state_dict['model'])
              self.opt.load_state_dict(state_dict['opt'])
  
      def _update_schedule(self):
          if self.args.lr_decay_rate < 1.:
              args = self.args
              t = self.updates
              base_ratio = args.min_lr / args.lr
              if t < args.lr_warmup_steps:
                  ratio = base_ratio + (1. - base_ratio) / max(1., args.lr_warmup_steps) * t
              else:
                  ratio = max(base_ratio, args.lr_decay_rate ** math.floor((t - args.lr_warmup_steps) /
                                                                           args.lr_decay_steps))
              self.opt.param_groups[0]['lr'] = args.lr * ratio
  
      def update(self, batch):
          self.network.train()
          self.opt.zero_grad()
          inputs, target = self.process_data(batch)
          output = self.network(inputs)
          summary = self.network.get_summary()
          loss = self.get_loss(output, target)
          loss.backward()
          grad_norm = torch.nn.utils.clip_grad_norm_(self.params, self.args.grad_clipping)
          assert grad_norm >= 0, 'encounter nan in gradients.'
          if isinstance(grad_norm, torch.Tensor):
              grad_norm = grad_norm.item()
          self.opt.step()
          self._update_schedule()
          self.updates += 1
          stats = {
              'updates': self.updates,
              'loss': loss.item(),
              'lr': self.opt.param_groups[0]['lr'],
              'gnorm': grad_norm,
              'summary': summary,
          }
          return stats
  
      def evaluate(self, data):
          self.network.eval()
          targets = []
          probabilities = []
          predictions = []
          losses = []
          for batch in tqdm(data[:self.args.eval_subset], desc='evaluating', leave=False):
              inputs, target = self.process_data(batch)
              with torch.no_grad():
                  output = self.network(inputs)
                  loss = self.get_loss(output, target)
                  pred = torch.argmax(output, dim=1)
                  prob = torch.nn.functional.softmax(output, dim=1)
                  losses.append(loss.item())
                  targets.extend(target.tolist())
                  probabilities.extend(prob.tolist())
                  predictions.extend(pred.tolist())
          outputs = {
              'target': targets,
              'prob': probabilities,
              'pred': predictions,
              'args': self.args,
          }
          stats = {
              'updates': self.updates,
              'loss': sum(losses[:-1]) / (len(losses) - 1) if len(losses) > 1 else sum(losses),
          }
          for metric in self.args.watch_metrics:
              if metric not in stats:  # multiple metrics could be computed by the same function
                  stats.update(metrics[metric](outputs))
          assert 'score' not in stats, 'metric name collides with "score"'
          eval_score = stats[self.args.metric]
          stats['score'] = eval_score
          return eval_score, stats  # first value is for early stopping
  
      def predict(self, batch):
          self.network.eval()
          inputs, _ = self.process_data(batch)
          with torch.no_grad():
              output = self.network(inputs)
              output = torch.nn.functional.softmax(output, dim=1)
          return output.tolist()
  
      def process_data(self, batch):
          text1 = torch.LongTensor(batch['text1']).to(self.device)
          text2 = torch.LongTensor(batch['text2']).to(self.device)
          mask1 = torch.ne(text1, self.args.padding).unsqueeze(2)
          mask2 = torch.ne(text2, self.args.padding).unsqueeze(2)
          inputs = {
              'text1': text1,
              'text2': text2,
              'mask1': mask1,
              'mask2': mask2,
          }
          if 'target' in batch:
              target = torch.LongTensor(batch['target']).to(self.device)
              return inputs, target
          return inputs, None
  
      @staticmethod
      def get_loss(logits, target):
          return f.cross_entropy(logits, target)
  
      def save(self, states, name=None):
          if name:
              filename = os.path.join(self.args.summary_dir, name)
          else:
              filename = os.path.join(self.args.summary_dir, f'{self.prefix}_{self.updates}.pt')
          params = {
              'state_dict': {
                  'model': self.network.state_dict(),
                  'opt': self.opt.state_dict(),
                  'updates': self.updates,
              },
              'args': self.args,
              'random_state': random.getstate(),
              'torch_state': torch.random.get_rng_state()
          }
          params.update(states)
          if self.args.cuda:
              params['torch_cuda_state'] = torch.cuda.get_rng_state()
          torch.save(params, filename)
  
      @classmethod
      def load(cls, file):
          checkpoint = torch.load(file, map_location=(
              lambda s, _: torch.serialization.default_restore_location(s, 'cpu')
          ))
          prev_args = checkpoint['args']
          # update args
          prev_args.output_dir = os.path.dirname(os.path.dirname(file))
          prev_args.summary_dir = os.path.join(prev_args.output_dir, prev_args.name)
          prev_args.cuda = prev_args.cuda and torch.cuda.is_available()
          return cls(prev_args, state_dict=checkpoint['state_dict']), checkpoint
  
      def num_parameters(self, exclude_embed=False):
          num_params = sum(p.numel() for p in self.network.parameters() if p.requires_grad)
          if exclude_embed:
              num_params -= 0 if self.args.fix_embeddings else next(self.network.embedding.parameters()).numel()
          return num_params
  
      def set_embeddings(self, embeddings):
          self.network.embedding.set_(embeddings)
  