class_docstr: Implements the beam search algorithm for decoding the most likely sequences.
  
  # Parameters
  
  end_index : `int`
      The index of the "stop" or "end" token in the target vocabulary.
  
  max_steps : `int`, optional (default = `50`)
      The maximum number of decoding steps to take, i.e. the maximum length
      of the predicted sequences.
  
  beam_size : `int`, optional (default = `10`)
      The width of the beam used.
  
  per_node_beam_size : `int`, optional (default = `beam_size`)
      The maximum number of candidates to consider per node, at each step in the search.
      If not given, this just defaults to `beam_size`. Setting this parameter
      to a number smaller than `beam_size` may give better results, as it can introduce
      more diversity into the search. See
      [*Beam Search Strategies for Neural Machine Translation*, Freitag and Al-Onaizan, 2017]
      (https://api.semanticscholar.org/CorpusID:2229477).
  
  sampler : `Sampler`, optional (default = `None`)
      An optional `Sampler` which is used to pick next candidate nodes and beams.
      If not specified, `DeterministicSampler` will be used, which just takes the
      `per_node_beam_size` most likely nodes and the `beam_size` most likely beams.
  
      Using the [`GumbelSampler`](#gumbelsampler), on the other hand, will give you
      [Stochastic Beam Search](https://api.semanticscholar.org/CorpusID:76662039).
  
  min_steps : `int`, optional (default = `None`)
      The minimum number of decoding steps to take, i.e. the minimum length of
      the predicted sequences. This does not include the start or end tokens. If `None`,
      no minimum is enforced.
  
  final_sequence_scorer : `FinalSequenceScorer`, optional (default = `None`)
      An optional `FinalSequenceScorer` which is used to score the final generated sequences.
      The output from this module is what is returned by the `search` method. If not
      specified, `SequenceLogProbabilityScorer` will be used, which scores the sequences
      by the sum of the token log probabilities.
  
  constraints: `List[Constraint]`, optional (default = `None`)
      An optional list of `Constraint`s which should be applied during beam search. If not
      provided, no constraints will be enforced.
  
  vocab: `Vocabulary`
      If `constraints` is not `None`, then `Vocabulary` will be passed to each constraint
      during its initialization. Having access to the vocabulary may be useful for certain
      contraints, e.g., to mask out invalid predictions during structured prediction.
  
      In a typical AllenNLP configuration file, this parameter does not get an entry under the
      "model", it gets specified as a top-level parameter, then is passed in to the model
      separately.
class_docstr_tok: ["Implements","the","beam","search","algorithm","for","decoding","the","most","likely","sequences",".","#","Parameters","end_index",":","`","int","`","The","index","of","the","``","stop","''","or","``","end","''","token","in","the","target","vocabulary",".","max_steps",":","`","int","`",",","optional","(","default","=","`","50","`",")","The","maximum","number","of","decoding","steps","to","take",",","i.e",".","the","maximum","length","of","the","predicted","sequences",".","beam_size",":","`","int","`",",","optional","(","default","=","`","10","`",")","The","width","of","the","beam","used",".","per_node_beam_size",":","`","int","`",",","optional","(","default","=","`","beam_size","`",")","The","maximum","number","of","candidates","to","consider","per","node",",","at","each","step","in","the","search",".","If","not","given",",","this","just","defaults","to","`","beam_size","`",".","Setting","this","parameter","to","a","number","smaller","than","`","beam_size","`","may","give","better","results",",","as","it","can","introduce","more","diversity","into","the","search",".","See","[","*","Beam","Search","Strategies","for","Neural","Machine","Translation","*",",","Freitag","and","Al-Onaizan",",","2017","]","(","https",":","\/\/api.semanticscholar.org\/CorpusID:2229477",")",".","sampler",":","`","Sampler","`",",","optional","(","default","=","`","None","`",")","An","optional","`","Sampler","`","which","is","used","to","pick","next","candidate","nodes","and","beams",".","If","not","specified",",","`","DeterministicSampler","`","will","be","used",",","which","just","takes","the","`","per_node_beam_size","`","most","likely","nodes","and","the","`","beam_size","`","most","likely","beams",".","Using","the","[","`","GumbelSampler","`","]","(","#","gumbelsampler",")",",","on","the","other","hand",",","will","give","you","[","Stochastic","Beam","Search","]","(","https",":","\/\/api.semanticscholar.org\/CorpusID:76662039",")",".","min_steps",":","`","int","`",",","optional","(","default","=","`","None","`",")","The","minimum","number","of","decoding","steps","to","take",",","i.e",".","the","minimum","length","of","the","predicted","sequences",".","This","does","not","include","the","start","or","end","tokens",".","If","`","None","`",",","no","minimum","is","enforced",".","final_sequence_scorer",":","`","FinalSequenceScorer","`",",","optional","(","default","=","`","None","`",")","An","optional","`","FinalSequenceScorer","`","which","is","used","to","score","the","final","generated","sequences",".","The","output","from","this","module","is","what","is","returned","by","the","`","search","`","method",".","If","not","specified",",","`","SequenceLogProbabilityScorer","`","will","be","used",",","which","scores","the","sequences","by","the","sum","of","the","token","log","probabilities",".","constraints",":","`","List","[","Constraint","]","`",",","optional","(","default","=","`","None","`",")","An","optional","list","of","`","Constraint","`","s","which","should","be","applied","during","beam","search",".","If","not","provided",",","no","constraints","will","be","enforced",".","vocab",":","`","Vocabulary","`","If","`","constraints","`","is","not","`","None","`",",","then","`","Vocabulary","`","will","be","passed","to","each","constraint","during","its","initialization",".","Having","access","to","the","vocabulary","may","be","useful","for","certain","contraints",",","e.g.",",","to","mask","out","invalid","predictions","during","structured","prediction",".","In","a","typical","AllenNLP","configuration","file",",","this","parameter","does","not","get","an","entry","under","the","``","model","''",",","it","gets","specified","as","a","top-level","parameter",",","then","is","passed","in","to","the","model","separately","."]
docstr: Given a starting state and a step function, apply beam search to find the
  most likely target sequences.
  
  !!! Note
      If your step function returns `-inf` for some log probabilities
      (like if you're using a masked log-softmax) then some of the "best"
      sequences returned may also have `-inf` log probability. Specifically
      this happens when the beam size is smaller than the number of actions
      with finite log probability (non-zero probability) returned by the step function.
      Therefore if you're using a mask you may want to check the results from `search`
      and potentially discard sequences with non-finite log probability.
  
  # Parameters
  
  start_predictions : `torch.Tensor`
      A tensor containing the initial predictions with shape `(batch_size,)`.
      Usually the initial predictions are just the index of the "start" token
      in the target vocabulary.
  
  start_state : `StateType`
      The initial state passed to the `step` function. Each value of the state dict
      should be a tensor of shape `(batch_size, *)`, where `*` means any other
      number of dimensions.
  
  step : `StepFunctionType`
      A function that is responsible for computing the next most likely tokens,
      given the current state and the predictions from the last time step.
      The function should accept two or three arguments:
  
      - a tensor of shape `(group_size,)` representing the index of the predicted
      tokens from the last time step,
      - the current state, a `StateType`, and
      - optionally, the timestep, an `int`.
  
      The `group_size` will be `batch_size * beam_size`, except in the initial
      step, for which it will just be `batch_size`.
  
      The function is expected to return a tuple, where the first element
      is a tensor of shape `(group_size, target_vocab_size)` containing
      the log probabilities of the tokens for the next step, and the second
      element is the updated state. The tensor in the state should have shape
      `(group_size, *)`, where `*` means any other number of dimensions.
  
  # Returns
  
  `Tuple[torch.Tensor, torch.Tensor]`
      Tuple of `(predictions, final_scores)`, where `predictions`
      has shape `(batch_size, beam_size, max_steps)` and `final_scores`
      has shape `(batch_size, beam_size)`.
docstr_tok: ["Given","a","starting","state","and","a","step","function",",","apply","beam","search","to","find","the","most","likely","target","sequences",".","!","!","!","Note","If","your","step","function","returns","`","-inf","`","for","some","log","probabilities","(","like","if","you","'re","using","a","masked","log-softmax",")","then","some","of","the","``","best","''","sequences","returned","may","also","have","`","-inf","`","log","probability",".","Specifically","this","happens","when","the","beam","size","is","smaller","than","the","number","of","actions","with","finite","log","probability","(","non-zero","probability",")","returned","by","the","step","function",".","Therefore","if","you","'re","using","a","mask","you","may","want","to","check","the","results","from","`","search","`","and","potentially","discard","sequences","with","non-finite","log","probability",".","#","Parameters","start_predictions",":","`","torch.Tensor","`","A","tensor","containing","the","initial","predictions","with","shape","`","(","batch_size",",",")","`",".","Usually","the","initial","predictions","are","just","the","index","of","the","``","start","''","token","in","the","target","vocabulary",".","start_state",":","`","StateType","`","The","initial","state","passed","to","the","`","step","`","function",".","Each","value","of","the","state","dict","should","be","a","tensor","of","shape","`","(","batch_size",",","*",")","`",",","where","`","*","`","means","any","other","number","of","dimensions",".","step",":","`","StepFunctionType","`","A","function","that","is","responsible","for","computing","the","next","most","likely","tokens",",","given","the","current","state","and","the","predictions","from","the","last","time","step",".","The","function","should","accept","two","or","three","arguments",":","-","a","tensor","of","shape","`","(","group_size",",",")","`","representing","the","index","of","the","predicted","tokens","from","the","last","time","step",",","-","the","current","state",",","a","`","StateType","`",",","and","-","optionally",",","the","timestep",",","an","`","int","`",".","The","`","group_size","`","will","be","`","batch_size","*","beam_size","`",",","except","in","the","initial","step",",","for","which","it","will","just","be","`","batch_size","`",".","The","function","is","expected","to","return","a","tuple",",","where","the","first","element","is","a","tensor","of","shape","`","(","group_size",",","target_vocab_size",")","`","containing","the","log","probabilities","of","the","tokens","for","the","next","step",",","and","the","second","element","is","the","updated","state",".","The","tensor","in","the","state","should","have","shape","`","(","group_size",",","*",")","`",",","where","`","*","`","means","any","other","number","of","dimensions",".","#","Returns","`","Tuple","[","torch.Tensor",",","torch.Tensor","]","`","Tuple","of","`","(","predictions",",","final_scores",")","`",",","where","`","predictions","`","has","shape","`","(","batch_size",",","beam_size",",","max_steps",")","`","and","`","final_scores","`","has","shape","`","(","batch_size",",","beam_size",")","`","."]
code: def search(
      self,
      start_predictions: torch.Tensor,
      start_state: StateType,
      step: StepFunctionType,
  ) -> Tuple[torch.Tensor, torch.Tensor]:
      """
      Given a starting state and a step function, apply beam search to find the
      most likely target sequences.
  
      !!! Note
          If your step function returns `-inf` for some log probabilities
          (like if you're using a masked log-softmax) then some of the "best"
          sequences returned may also have `-inf` log probability. Specifically
          this happens when the beam size is smaller than the number of actions
          with finite log probability (non-zero probability) returned by the step function.
          Therefore if you're using a mask you may want to check the results from `search`
          and potentially discard sequences with non-finite log probability.
  
      # Parameters
  
      start_predictions : `torch.Tensor`
          A tensor containing the initial predictions with shape `(batch_size,)`.
          Usually the initial predictions are just the index of the "start" token
          in the target vocabulary.
  
      start_state : `StateType`
          The initial state passed to the `step` function. Each value of the state dict
          should be a tensor of shape `(batch_size, *)`, where `*` means any other
          number of dimensions.
  
      step : `StepFunctionType`
          A function that is responsible for computing the next most likely tokens,
          given the current state and the predictions from the last time step.
          The function should accept two or three arguments:
  
          - a tensor of shape `(group_size,)` representing the index of the predicted
          tokens from the last time step,
          - the current state, a `StateType`, and
          - optionally, the timestep, an `int`.
  
          The `group_size` will be `batch_size * beam_size`, except in the initial
          step, for which it will just be `batch_size`.
  
          The function is expected to return a tuple, where the first element
          is a tensor of shape `(group_size, target_vocab_size)` containing
          the log probabilities of the tokens for the next step, and the second
          element is the updated state. The tensor in the state should have shape
          `(group_size, *)`, where `*` means any other number of dimensions.
  
      # Returns
  
      `Tuple[torch.Tensor, torch.Tensor]`
          Tuple of `(predictions, final_scores)`, where `predictions`
          has shape `(batch_size, beam_size, max_steps)` and `final_scores`
          has shape `(batch_size, beam_size)`.
      """
      step_signature = signature(step)
      if len(step_signature.parameters) < 3:
          # If the step function we're given does not take the time step argument, wrap it
          # in one that does.
          old_step = cast(StepFunctionTypeNoTimestep, step)
  
          def new_step(
              last_predictions: torch.Tensor,
              state: Dict[str, torch.Tensor],
              time_step: int,
          ):
              return old_step(last_predictions, state)
  
          return self._search(start_predictions, start_state, new_step)
      else:
          return self._search(
              start_predictions,
              start_state,
              cast(StepFunctionTypeWithTimestep, step),
          )
  
imports: from inspect import signature
  from typing import Any, List, Callable, Tuple, Dict, cast, TypeVar, Optional
  import copy
  import warnings
  import torch
  from allennlp.common import Lazy, Registrable
  from allennlp.common.checks import ConfigurationError
  from allennlp.data import Vocabulary
  from allennlp.nn.util import min_value_of_dtype
label_desc: Use image node_id 3 for calling the BeamSearch obj's underlying member method code with example usage: obj.search(start_predictions, start_state, step) and returns: self, self, old_step