code: def _search_beam(
      self,
      beam_idx: int,
      batch_offset: int,
      lprobs: Tensor,
      step_scores: Tensor,
  ) -> Optional[BeamStep]:
      # Ignore the generated indices for the prompt sequences.
      if self.step_nr < self.max_prompt_len:
          assert self.prompt_mask is not None
  
          # Check if the current beam is in a prompt sequence.
          if self.prompt_mask[batch_offset, self.step_nr]:
              # The size of a beam in a prompt sequence must be always 1.
              assert len(lprobs) == 1
  
              seq_index = torch.tensor(
                  [batch_offset], device=lprobs.device
              )
  
              # We just extract the prompt step along with its score and treat
              # it as the next beam step. So we keep a beam of size 1 until we
              # reach the end of the prompt.
              vocab_index = self.seqs[
                  batch_offset, self.step_nr : self.step_nr + 1
              ]
  
              score = (
                  step_scores[0, self.step_nr - 1]
                  + lprobs[0, vocab_index]
              )
  
              return BeamStep(seq_index, vocab_index, score)
      else:
          self.prompt_mask = None  # Not needed anymore, release.
  
      # We use the same beam search method as in fairseq, where we take the
      # best 2 x `beam_size` candidates and choose the first `beam_size` of
      # these which don't predict EOS to continue with.
      # (2 x B)
      next_step = self.algorithm(
          self.beam_size, lprobs, step_scores[:, : self.step_nr]
      )
  
      # Translate the sequence indices from beam to batch.
      next_step.seq_indices += batch_offset
  
      # (2 x B)
      eos_mask = next_step.vocab_indices == self.eos_idx
  
      # Consider EOS only when it's among the top `beam_size` indices.
      # (F)
      eos_seq_indices = next_step.seq_indices[
          : self.beam_size
      ].masked_select(eos_mask[: self.beam_size])
  
      # If one or more sequences have reached EOS, move them to the output.
      if len(eos_seq_indices) > 0:
          # (F)
          eos_scores = next_step.scores[: self.beam_size].masked_select(
              eos_mask[: self.beam_size]
          )
  
          for seq_idx, score in zip(eos_seq_indices, eos_scores):
              # If `True`, it means we have found `beam_size` hypotheses for
              # this beam.
              if self._finish_sequence(int(seq_idx), score):
                  return None
  
          # Filter out the sequences that have reached EOS.
          seq_mask = ~eos_mask
  
          next_step = next_step.masked_select(seq_mask)
  
      # We can have at most `beam_size` sequences in the beam.
      return next_step.first(self.beam_size)
  
imports: from __future__ import annotations
  from abc import ABC, abstractmethod
  from dataclasses import dataclass
  from typing import Dict, List, Optional, Sequence, Tuple, Union, final
  import torch
  from torch import Tensor
  from torch.nn.functional import log_softmax
  from fairseq2.data import VocabularyInfo
  from fairseq2.generation.generator import Hypothesis, Seq2SeqGenerator, Seq2SeqGeneratorOutput, SequenceGenerator, SequenceGeneratorOutput, StepHook
  from fairseq2.generation.step_processor import StepProcessor
  from fairseq2.models.decoder import DecoderModel
  from fairseq2.models.encoder_decoder import EncoderDecoderModel
  from fairseq2.models.sequence import SequenceModelOutput
  from fairseq2.nn.incremental_state import IncrementalStateBag
  from fairseq2.nn.padding import PaddingMask
  from fairseq2.typing import finaloverride, override
label_desc: Use image node_id 6 for calling the _BeamSearchSequenceGeneratorOpBase obj's underlying member method code with example usage: obj._search_beam(beam_idx, batch_offset, lprobs, step_scores) and returns: next_step, BeamStep, None