code: def test_search(self):
      # The simple transition system starts at some number, adds one or two at each state, and
      # tries to get to 4.  The highest scoring path has the shortest length and the highest
      # numbers (so always add two, unless you're at 3).  From -3, there are lots of possible
      # sequences: [-2, -1, 0, 1, 2, 3, 4], [-1, 1, 3, 4], ...  We'll specify a few of those up
      # front as "allowed", and use that to test the constrained beam search implementation.
      initial_state = SimpleDecoderState(
          [0], [[]], [Variable(torch.Tensor([0.0]))], [-3]
      )
      beam_size = 3
      # pylint: disable=bad-whitespace
      allowed_sequences = Variable(
          torch.Tensor(
              [
                  [
                      [-2, -1, 0, 1, 2, 3, 4],
                      [-2, 0, 2, 4, -1, -1, -1],
                      [-1, 1, 3, 4, -1, -1, -1],
                      [-2, -1, 0, 1, 2, 4, -1],
                      [-1, 0, 1, 2, 3, 4, -1],
                      [-1, 1, 2, 3, 4, -1, -1],
                  ]
              ]
          )
      )
      # pylint: enable=bad-whitespace
      mask = Variable(
          torch.Tensor(
              [
                  [
                      [1, 1, 1, 1, 1, 1, 1],
                      [1, 1, 1, 1, 0, 0, 0],
                      [1, 1, 1, 1, 0, 0, 0],
                      [1, 1, 1, 1, 1, 1, 0],
                      [1, 1, 1, 1, 1, 1, 0],
                      [1, 1, 1, 1, 1, 0, 0],
                  ]
              ]
          )
      )
  
      beam_search = ConstrainedBeamSearch(
          beam_size, allowed_sequences, mask
      )
  
      # Including the value in the score will make us pick states that have higher numbers first.
      # So with a beam size of 3, we'll get all of the states that start with `-1` after the
      # first step, even though in the end one of the states that starts with `-2` is better than
      # two of the states that start with `-1`.
      decoder_step = SimpleDecoderStep(include_value_in_score=True)
      best_states = beam_search.search(initial_state, decoder_step)
  
      assert len(best_states) == 1
      assert best_states[0][0].action_history[0] == [-1, 1, 3, 4]
      assert best_states[0][1].action_history[0] == [-1, 1, 2, 3, 4]
      assert best_states[0][2].action_history[0] == [-1, 0, 1, 2, 3, 4]
  
      # With a beam size of 6, we should get the other allowed path of length 4 as the second
      # best result.
      beam_size = 6
      beam_search = ConstrainedBeamSearch(
          beam_size, allowed_sequences, mask
      )
      decoder_step = SimpleDecoderStep(include_value_in_score=True)
      best_states = beam_search.search(initial_state, decoder_step)
  
      assert len(best_states) == 1
      assert best_states[0][0].action_history[0] == [-1, 1, 3, 4]
      assert best_states[0][1].action_history[0] == [-2, 0, 2, 4]
  
file_path: magnitude/pymagnitude/third_party/allennlp/tests/nn/decoding/constrained_beam_search_test.py
imports: from __future__ import absolute_import
  import torch
  from torch.autograd import Variable
  from allennlp.common.testing import AllenNlpTestCase
  from allennlp.nn.decoding import ConstrainedBeamSearch
  from .simple_transition_system import SimpleDecoderState, SimpleDecoderStep
label_desc: Use image node_id 1 for calling the TestConstrainedBeamSearch obj's underlying member method code with example usage: obj.test_search() without return types