class_docstr: Diverse Beam Search.
  
  See "Diverse Beam Search: Decoding Diverse Solutions from Neural Sequence
  Models" for details.
  
  We implement cumulative diversity penalty here as default, optionally provide Hamming diversity described
  in the original paper, and a way to interpolate between the two through diversity_discount.
  
  Take the example below for illustration of cumulative diversity implemented.
  A) I like dogs.
  B) I like ____.
  C) There are ___.
  And we are at step=2, trying to fill in the blank:
  
  Hamming diversity:
  Penalty for B from A is 1 for "dogs" and 0 for any other words like "cats".
  Penalty for C from A is 1 for "dogs" and 0 for any other words like "cats".
  
  Cumulative diversity (default):
  Penalty for B from A is 3 for "dogs" and 0 for any other words like "cats".
  Penalty for C from A is 1 for "dogs" and 0 for any other words like "cats".
  B and C differ because B matches with A for "I" and "like" at respective steps incurring 2 cumulative penalty.
  
  Using divesrity_discount to interpolate between the two:
  if diverstiy_discount = 0.5, then
  Penalty for B from A is 1.75 (1 + 0.5 + 0.25) for "dogs" and 0 for any other words like "cats".
  Penalty for C from A is 1 for "dogs" and 0 for any other words like "cats".
  "I" and "like" matched for B and A at step 0 and 1 respectively. Since "I" is two steps away and "like" is one step away, they are discounted by (0.5)^2 and 0.5 respectively.
  When diversity_discount = 0, we recover Hammning diversity and when diversity_discount = 1, we recover cumulative diversity.
  
  NB: During beam search for each diversity group, `candidate_mutiple` is set to 1 rather than BeamSearch default(2).
  This is to ensure we have final `beam_size` candidates so that no diversity groups would be dropped during final token selection in sequence generation.
  For full backwards compatibility, use diversity_discount=0 and candidate_multiple=2.
class_docstr_tok: ["Diverse","Beam","Search",".","See","``","Diverse","Beam","Search",":","Decoding","Diverse","Solutions","from","Neural","Sequence","Models","''","for","details",".","We","implement","cumulative","diversity","penalty","here","as","default",",","optionally","provide","Hamming","diversity","described","in","the","original","paper",",","and","a","way","to","interpolate","between","the","two","through","diversity_discount",".","Take","the","example","below","for","illustration","of","cumulative","diversity","implemented",".","A",")","I","like","dogs",".","B",")","I","like","____",".","C",")","There","are","___",".","And","we","are","at","step=2",",","trying","to","fill","in","the","blank",":","Hamming","diversity",":","Penalty","for","B","from","A","is","1","for","``","dogs","''","and","0","for","any","other","words","like","``","cats","''",".","Penalty","for","C","from","A","is","1","for","``","dogs","''","and","0","for","any","other","words","like","``","cats","''",".","Cumulative","diversity","(","default",")",":","Penalty","for","B","from","A","is","3","for","``","dogs","''","and","0","for","any","other","words","like","``","cats","''",".","Penalty","for","C","from","A","is","1","for","``","dogs","''","and","0","for","any","other","words","like","``","cats","''",".","B","and","C","differ","because","B","matches","with","A","for","``","I","''","and","``","like","''","at","respective","steps","incurring","2","cumulative","penalty",".","Using","divesrity_discount","to","interpolate","between","the","two",":","if","diverstiy_discount","=","0.5",",","then","Penalty","for","B","from","A","is","1.75","(","1","+","0.5","+","0.25",")","for","``","dogs","''","and","0","for","any","other","words","like","``","cats","''",".","Penalty","for","C","from","A","is","1","for","``","dogs","''","and","0","for","any","other","words","like","``","cats","''",".","``","I","''","and","``","like","''","matched","for","B","and","A","at","step","0","and","1","respectively",".","Since","``","I","''","is","two","steps","away","and","``","like","''","is","one","step","away",",","they","are","discounted","by","(","0.5",")","^2","and","0.5","respectively",".","When","diversity_discount","=","0",",","we","recover","Hammning","diversity","and","when","diversity_discount","=","1",",","we","recover","cumulative","diversity",".","NB",":","During","beam","search","for","each","diversity","group",",","`","candidate_mutiple","`","is","set","to","1","rather","than","BeamSearch","default","(","2",")",".","This","is","to","ensure","we","have","final","`","beam_size","`","candidates","so","that","no","diversity","groups","would","be","dropped","during","final","token","selection","in","sequence","generation",".","For","full","backwards","compatibility",",","use","diversity_discount=0","and","candidate_multiple=2","."]
code: def step(
      self,
      step: int,
      lprobs,
      scores,
      prev_output_tokens: Optional[Tensor] = None,
      original_batch_idxs: Optional[Tensor] = None,
  ):
      bsz, beam_size, vocab_size = lprobs.size()
      if beam_size % self.num_groups != 0:
          raise ValueError(
              "DiverseBeamSearch requires --beam to be divisible by the number of groups"
          )
  
      # initialize diversity penalty
      diversity_buf = torch.zeros(lprobs[:, 0, :].size()).to(lprobs)
  
      scores_G, beams_G = [], []
  
      # pre-allocating tensor for indices for all groups
      indices_G_stacked = torch.empty(
          bsz,
          int(beam_size / self.num_groups) * self.candidate_multiple,
          self.num_groups,
          dtype=torch.long,
          device=lprobs.device,
      )
  
      for g in range(self.num_groups):
          lprobs_g = lprobs[:, g :: self.num_groups, :]
          scores_g = (
              scores[:, g :: self.num_groups, :] if step > 0 else None
          )
  
          diversity_buf.zero_()
          # apply diversity penalty
          if g > 0:
              indices_ = indices_G_stacked[:, :, :g]
              if step > 0:
                  penalty_val = (
                      1 + self.group_overlap[original_batch_idxs, g, :g]
                  )
                  penalty_val = penalty_val.unsqueeze(1)
              else:
                  penalty_val = torch.ones(bsz, 1, 1)
              diversity_buf.scatter_add_(
                  1,
                  indices_.reshape(bsz, -1),
                  penalty_val.expand(indices_.size())
                  .reshape(bsz, -1)
                  .to(diversity_buf),
              )
  
              lprobs_g = torch.add(
                  lprobs_g,
                  other=diversity_buf.unsqueeze(1),
                  alpha=self.diversity_strength,
              )
          else:
              lprobs_g = lprobs_g.contiguous()
  
          scores_buf, indices_buf, beams_buf = self.beam.step(
              step,
              lprobs_g,
              scores_g,
              candidate_multiple=self.candidate_multiple,
          )
          beams_buf.mul_(self.num_groups).add_(g)
  
          scores_G.append(scores_buf.clone())
          beams_G.append(beams_buf.clone())
  
          indices_G_stacked[:, :, g] = indices_buf
  
      # interleave results from different groups
      scores_buf = torch.stack(scores_G, dim=2).view(bsz, -1)
      indices_buf = indices_G_stacked.view(bsz, -1)
      beams_buf = torch.stack(beams_G, dim=2).view(bsz, -1)
      # find num of overlapped tokens for each group pair
      # then discount it for next timestamp
      overlap = self.diversity_discount * torch.sum(
          indices_G_stacked.unsqueeze(2).eq(
              indices_G_stacked.unsqueeze(3)
          ),
          dim=1,
      )
      if step == 0:
          self.group_overlap = overlap
      else:
          self.group_overlap[original_batch_idxs] = (
              self.group_overlap[original_batch_idxs]
              * self.diversity_discount
              + overlap
          )
  
      return scores_buf, indices_buf, beams_buf
  
imports: import math
  from typing import List, Optional
  import torch
  import torch.nn
  from fairseq.token_generation_constraints import ConstraintState, OrderedConstraintState, UnorderedConstraintState
  from torch import Tensor
label_desc: Use image node_id 2 for calling the DiverseBeamSearch obj's underlying member method code with example usage: obj.step(step, lprobs, scores, prev_output_tokens, original_batch_idxs) and returns: scores_buf,  indices_buf,  beams_buf