class_docstr: This class implements beam search over transition sequences given an initial ``DecoderState``,
  a ``DecoderStep``, and a list of allowed transition sequences.  We will do a beam search `over
  the list of allowed sequences` and return the highest scoring states found by the beam.  This
  is only actually a `beam search` if your beam size is smaller than the list of allowed
  transition sequences; otherwise, we are just scoring and sorting the sequences using a prefix
  tree.
  
  The initial ``DecoderState`` is assumed to be `batched`.  The value we return from the search
  is a dictionary from batch indices to ranked finished states.
  
  IMPORTANT: We assume that the ``DecoderStep`` that you are using returns possible next states
  in sorted order, so we do not do an additional sort inside of
  ``ConstrainedBeamSearch.search()``.  If you're implementing your own ``DecoderStep``, you must
  ensure that you've sorted the states that you return.
  
  Parameters
  ----------
  beam_size : ``int``
      The beam size to use.
  allowed_sequences : ``torch.Tensor``
      A ``(batch_size, num_sequences, sequence_length)`` tensor containing the transition
      sequences that we will search in.  The values in this tensor must match whatever the
      ``DecoderState`` keeps in its ``action_history`` variable (typically this is action
      indices).
  allowed_sequence_mask : ``torch.Tensor``
      A ``(batch_size, num_sequences, sequence_length)`` tensor indicating whether each entry in
      the ``allowed_sequences`` tensor is padding.  The allowed sequences could be padded both on
      the ``num_sequences`` dimension and the ``sequence_length`` dimension.
docstr: Parameters
  ----------
  initial_state : ``DecoderState``
      The starting state of our search.  This is assumed to be `batched`, and our beam search
      is batch-aware - we'll keep ``beam_size`` states around for each instance in the batch.
  decoder_step : ``DecoderStep``
      The ``DecoderStep`` object that defines and scores transitions from one state to the
      next.
  
  Returns
  -------
  best_states : ``Dict[int, List[DecoderState]]``
      This is a mapping from batch index to the top states for that instance.
docstr_tok: ["Parameters","--","--","--","--","--","initial_state",":","``","DecoderState","``","The","starting","state","of","our","search",".","This","is","assumed","to","be","`","batched","`",",","and","our","beam","search","is","batch-aware","-","we","'ll","keep","``","beam_size","``","states","around","for","each","instance","in","the","batch",".","decoder_step",":","``","DecoderStep","``","The","``","DecoderStep","``","object","that","defines","and","scores","transitions","from","one","state","to","the","next",".","Returns","--","--","--","-","best_states",":","``","Dict","[","int",",","List","[","DecoderState","]","]","``","This","is","a","mapping","from","batch","index","to","the","top","states","for","that","instance","."]
code: def search(self, initial_state, decoder_step):
      """
      Parameters
      ----------
      initial_state : ``DecoderState``
          The starting state of our search.  This is assumed to be `batched`, and our beam search
          is batch-aware - we'll keep ``beam_size`` states around for each instance in the batch.
      decoder_step : ``DecoderStep``
          The ``DecoderStep`` object that defines and scores transitions from one state to the
          next.
  
      Returns
      -------
      best_states : ``Dict[int, List[DecoderState]]``
          This is a mapping from batch index to the top states for that instance.
      """
      finished_states = defaultdict(list)
      states = [initial_state]
      step_num = 0
      while states:
          step_num += 1
          next_states = defaultdict(list)
          grouped_state = states[0].combine_states(states)
          allowed_actions = []
          for batch_index, action_history in izip(
              grouped_state.batch_indices, grouped_state.action_history
          ):
              allowed_actions.append(
                  self._allowed_transitions[batch_index][
                      tuple(action_history)
                  ]
              )
          for next_state in decoder_step.take_step(
              grouped_state,
              max_actions=self._beam_size,
              allowed_actions=allowed_actions,
          ):
              # NOTE: we're doing state.batch_indices[0] here (and similar things below),
              # hard-coding a group size of 1.  But, our use of `next_state.is_finished()`
              # already checks for that, as it crashes if the group size is not 1.
              batch_index = next_state.batch_indices[0]
              if next_state.is_finished():
                  finished_states[batch_index].append(next_state)
              else:
                  next_states[batch_index].append(next_state)
          states = []
          for batch_index, batch_states in list(next_states.items()):
              # The states from the generator are already sorted, so we can just take the first
              # ones here, without an additional sort.
              states.extend(batch_states[: self._beam_size])
      best_states = {}
      for batch_index, batch_states in list(finished_states.items()):
          # The time this sort takes is pretty negligible, no particular need to optimize this
          # yet.  Maybe with a larger beam size...
          finished_to_sort = [
              (-state.score[0].item(), state) for state in batch_states
          ]
          finished_to_sort.sort(key=lambda x: x[0])
          best_states[batch_index] = [
              state[1] for state in finished_to_sort[: self._beam_size]
          ]
      return best_states
  
file_path: magnitude/pymagnitude/third_party/allennlp/nn/decoding/constrained_beam_search.py
imports: from __future__ import absolute_import
  from collections import defaultdict
  import torch
  from allennlp.nn.decoding import util
  from allennlp.nn.decoding.decoder_step import DecoderStep
  from allennlp.nn.decoding.decoder_state import DecoderState
label_desc: Use image node_id 2 for calling the ConstrainedBeamSearch obj's underlying member method code with example usage: obj.search(initial_state, decoder_step) and returns: best_states