code: def test_search(self):
      beam_search = BeamSearch.from_params(Params({"beam_size": 4}))
      initial_state = SimpleDecoderState(
          [0, 1, 2, 3],
          [[], [], [], []],
          [
              torch.Tensor([0.0]),
              torch.Tensor([0.0]),
              torch.Tensor([0.0]),
              torch.Tensor([0.0]),
          ],
          [-3, 1, -20, 5],
      )
      decoder_step = SimpleDecoderStep(include_value_in_score=True)
      best_states = beam_search.search(
          5,
          initial_state,
          decoder_step,
          keep_final_unfinished_states=False,
      )
  
      # Instance with batch index 2 needed too many steps to finish, and batch index 3 had no
      # path to get to a finished state.  (See the simple transition system definition; goal is
      # to end up at 4, actions are either add one or two to starting value.)
      assert len(best_states) == 2
      assert best_states[0][0].action_history[0] == [-1, 1, 3, 4]
      assert best_states[1][0].action_history[0] == [3, 4]
  
      best_states = beam_search.search(
          5,
          initial_state,
          decoder_step,
          keep_final_unfinished_states=True,
      )
  
      # Now we're keeping final unfinished states, which allows a "best state" for the instances
      # that didn't have one before.  Our previous best states for the instances that finish
      # doesn't change, because the score for taking another step is always negative at these
      # values.
      assert len(best_states) == 4
      assert best_states[0][0].action_history[0] == [-1, 1, 3, 4]
      assert best_states[1][0].action_history[0] == [3, 4]
      assert best_states[2][0].action_history[0] == [
          -18,
          -16,
          -14,
          -12,
          -10,
      ]
      assert best_states[3][0].action_history[0] == [7, 9, 11, 13, 15]
  
file_path: magnitude/pymagnitude/third_party/allennlp/tests/nn/decoding/beam_search_test.py
imports: from __future__ import absolute_import
  import torch
  from allennlp.common import Params
  from allennlp.common.testing import AllenNlpTestCase
  from allennlp.nn.decoding import BeamSearch
  from .simple_transition_system import SimpleDecoderState, SimpleDecoderStep
label_desc: Use image node_id 1 for calling the TestBeamSearch obj's underlying member method code with example usage: obj.test_search() without return types