class_docstr: This class implements a trainer that minimizes the expected value of a cost function over the
  space of some candidate sequences produced by a decoder. We generate the candidate sequences by
  performing beam search (which is one of the two popular ways of getting these sequences, the
  other one being sampling; see "Classical Structured Prediction Losses for Sequence to Sequence
  Learning" by Edunov et al., 2017 for more details).
  Note that we do not have a notion of targets here, so we're breaking the API of DecoderTrainer
  a bit.
  
  Parameters
  ----------
  beam_size : ``int``
  noramlize_by_length : ``bool``
      Should the log probabilities be normalized by length before renormalizing them? Edunov et
      al. do this in their work.
  max_decoding_steps : ``int``
      The maximum number of steps we should take during decoding.
  max_num_decoded_sequences : ``int``, optional (default=1)
      Maximum number of sorted decoded sequences to return. Defaults to 1.
  max_num_finished_states : ``int``, optional (default = None)
      Maximum number of finished states to keep after search. This is to finished states as
      ``beam_size`` is to unfinished ones. Costs are computed for only these number of states per
      instance. If not set, we will keep all the finished states.
class_docstr_tok: ["This","class","implements","a","trainer","that","minimizes","the","expected","value","of","a","cost","function","over","the","space","of","some","candidate","sequences","produced","by","a","decoder",".","We","generate","the","candidate","sequences","by","performing","beam","search","(","which","is","one","of","the","two","popular","ways","of","getting","these","sequences",",","the","other","one","being","sampling",";","see","``","Classical","Structured","Prediction","Losses","for","Sequence","to","Sequence","Learning","''","by","Edunov","et","al.",",","2017","for","more","details",")",".","Note","that","we","do","not","have","a","notion","of","targets","here",",","so","we","'re","breaking","the","API","of","DecoderTrainer","a","bit",".","Parameters","--","--","--","--","--","beam_size",":","``","int","``","noramlize_by_length",":","``","bool","``","Should","the","log","probabilities","be","normalized","by","length","before","renormalizing","them","?","Edunov","et","al",".","do","this","in","their","work",".","max_decoding_steps",":","``","int","``","The","maximum","number","of","steps","we","should","take","during","decoding",".","max_num_decoded_sequences",":","``","int","``",",","optional","(","default=1",")","Maximum","number","of","sorted","decoded","sequences","to","return",".","Defaults","to","1.","max_num_finished_states",":","``","int","``",",","optional","(","default","=","None",")","Maximum","number","of","finished","states","to","keep","after","search",".","This","is","to","finished","states","as","``","beam_size","``","is","to","unfinished","ones",".","Costs","are","computed","for","only","these","number","of","states","per","instance",".","If","not","set",",","we","will","keep","all","the","finished","states","."]
docstr: This method can be used to prune the set of unfinished states on a beam or finished states
  at the end of search. In the former case, the states need not be sorted because the all come
  from the same decoding step, which does the sorting. However, if the states are finished and
  this method is called at the end of the search, they need to be sorted because they come
  from different decoding steps.
docstr_tok: ["This","method","can","be","used","to","prune","the","set","of","unfinished","states","on","a","beam","or","finished","states","at","the","end","of","search",".","In","the","former","case",",","the","states","need","not","be","sorted","because","the","all","come","from","the","same","decoding","step",",","which","does","the","sorting",".","However",",","if","the","states","are","finished","and","this","method","is","called","at","the","end","of","the","search",",","they","need","to","be","sorted","because","they","come","from","different","decoding","steps","."]
code: def _prune_beam(states, beam_size, sort_states=False):
      """
      This method can be used to prune the set of unfinished states on a beam or finished states
      at the end of search. In the former case, the states need not be sorted because the all come
      from the same decoding step, which does the sorting. However, if the states are finished and
      this method is called at the end of the search, they need to be sorted because they come
      from different decoding steps.
      """
      states_by_batch_index = defaultdict(list)
      for state in states:
          assert len(state.batch_indices) == 1
          batch_index = state.batch_indices[0]
          states_by_batch_index[batch_index].append(state)
      pruned_states = []
      for _, instance_states in list(states_by_batch_index.items()):
          if sort_states:
              scores = torch.cat(
                  [state.score[0].view(-1) for state in instance_states]
              )
              _, sorted_indices = scores.sort(-1, descending=True)
              sorted_states = [
                  instance_states[i]
                  for i in sorted_indices.detach().cpu().numpy()
              ]
              instance_states = sorted_states
          for state in instance_states[:beam_size]:
              pruned_states.append(state)
      return pruned_states
  
file_path: magnitude/pymagnitude/third_party/allennlp/nn/decoding/decoder_trainers/expected_risk_minimization.py
imports: from __future__ import division
  from __future__ import absolute_import
  from collections import defaultdict
  import torch
  from allennlp.nn.decoding.decoder_step import DecoderStep
  from allennlp.nn.decoding.decoder_state import DecoderState
  from allennlp.nn.decoding.decoder_trainers.decoder_trainer import DecoderTrainer
  from allennlp.nn import util
label_desc: Use image node_id 4 for calling the ExpectedRiskMinimization obj's underlying member method code with example usage: obj._prune_beam(states, beam_size, sort_states) and returns: pruned_states