text: import unittest
  from onmt.modules.copy_generator import CopyGenerator, CopyGeneratorLoss
  
  import itertools
  from copy import deepcopy
  
  import torch
  from torch.nn.functional import softmax
  
  from onmt.tests.utils_for_tests import product_dict
  
  
  class TestCopyGenerator(unittest.TestCase):
      INIT_CASES = list(product_dict(
          input_size=[172],
          output_size=[319],
          pad_idx=[0, 39],
      ))
      PARAMS = list(product_dict(
          batch_size=[1, 14],
          max_seq_len=[23],
          tgt_max_len=[50],
          n_extra_words=[107]
      ))
  
      @classmethod
      def dummy_inputs(cls, params, init_case):
          hidden = torch.randn((params["batch_size"] * params["tgt_max_len"],
                                init_case["input_size"]))
          attn = torch.randn((params["batch_size"] * params["tgt_max_len"],
                              params["max_seq_len"]))
          src_map = torch.randn((params["max_seq_len"], params["batch_size"],
                                 params["n_extra_words"]))
          return hidden, attn, src_map
  
      @classmethod
      def expected_shape(cls, params, init_case):
          return params["tgt_max_len"] * params["batch_size"], \
                 init_case["output_size"] + params["n_extra_words"]
  
      def test_copy_gen_forward_shape(self):
          for params, init_case in itertools.product(
                  self.PARAMS, self.INIT_CASES):
              cgen = CopyGenerator(**init_case)
              dummy_in = self.dummy_inputs(params, init_case)
              res = cgen(*dummy_in)
              expected_shape = self.expected_shape(params, init_case)
              self.assertEqual(res.shape, expected_shape, init_case.__str__())
  
      def test_copy_gen_outp_has_no_prob_of_pad(self):
          for params, init_case in itertools.product(
                  self.PARAMS, self.INIT_CASES):
              cgen = CopyGenerator(**init_case)
              dummy_in = self.dummy_inputs(params, init_case)
              res = cgen(*dummy_in)
              self.assertTrue(
                  res[:, init_case["pad_idx"]].allclose(torch.tensor(0.0)))
  
      def test_copy_gen_trainable_params_update(self):
          for params, init_case in itertools.product(
                  self.PARAMS, self.INIT_CASES):
              cgen = CopyGenerator(**init_case)
              trainable_params = {n: p for n, p in cgen.named_parameters()
                                  if p.requires_grad}
              assert len(trainable_params) > 0  # sanity check
              old_weights = deepcopy(trainable_params)
              dummy_in = self.dummy_inputs(params, init_case)
              res = cgen(*dummy_in)
              pretend_loss = res.sum()
              pretend_loss.backward()
              dummy_optim = torch.optim.SGD(trainable_params.values(), 1)
              dummy_optim.step()
              for param_name in old_weights.keys():
                  self.assertTrue(
                      trainable_params[param_name]
                      .ne(old_weights[param_name]).any(),
                      param_name + " " + init_case.__str__())
  
  
  class TestCopyGeneratorLoss(unittest.TestCase):
      INIT_CASES = list(product_dict(
          vocab_size=[172],
          unk_index=[0, 39],
          ignore_index=[1, 17],  # pad idx
          force_copy=[True, False]
      ))
      PARAMS = list(product_dict(
          batch_size=[1, 14],
          tgt_max_len=[50],
          n_extra_words=[107]
      ))
  
      @classmethod
      def dummy_inputs(cls, params, init_case):
          n_unique_src_words = 13
          scores = torch.randn((params["batch_size"] * params["tgt_max_len"],
                                init_case["vocab_size"] + n_unique_src_words))
          scores = softmax(scores, dim=1)
          align = torch.randint(0, n_unique_src_words,
                                (params["batch_size"] * params["tgt_max_len"],))
          target = torch.randint(0, init_case["vocab_size"],
                                 (params["batch_size"] * params["tgt_max_len"],))
          target[0] = init_case["unk_index"]
          target[1] = init_case["ignore_index"]
          return scores, align, target
  
      @classmethod
      def expected_shape(cls, params, init_case):
          return (params["batch_size"] * params["tgt_max_len"],)
  
      def test_copy_loss_forward_shape(self):
          for params, init_case in itertools.product(
                  self.PARAMS, self.INIT_CASES):
              loss = CopyGeneratorLoss(**init_case)
              dummy_in = self.dummy_inputs(params, init_case)
              res = loss(*dummy_in)
              expected_shape = self.expected_shape(params, init_case)
              self.assertEqual(res.shape, expected_shape, init_case.__str__())
  
      def test_copy_loss_ignore_index_is_ignored(self):
          for params, init_case in itertools.product(
                  self.PARAMS, self.INIT_CASES):
              loss = CopyGeneratorLoss(**init_case)
              scores, align, target = self.dummy_inputs(params, init_case)
              res = loss(scores, align, target)
              should_be_ignored = (target == init_case["ignore_index"]).nonzero(
                  as_tuple=False)
              assert len(should_be_ignored) > 0  # otherwise not testing anything
              self.assertTrue(res[should_be_ignored].allclose(torch.tensor(0.0)))
  
      def test_copy_loss_output_range_is_positive(self):
          for params, init_case in itertools.product(
                  self.PARAMS, self.INIT_CASES):
              loss = CopyGeneratorLoss(**init_case)
              dummy_in = self.dummy_inputs(params, init_case)
              res = loss(*dummy_in)
              self.assertTrue((res >= 0).all())
  