text: from dataclasses import dataclass
  from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, Union
  
  import torch
  from torch.nn import functional as F
  
  from transformers.file_utils import ModelOutput
  from transformers.generation_beam_search import BeamScorer, BeamSearchScorer
  from transformers.generation_logits_process import (
      EncoderNoRepeatNGramLogitsProcessor,
      ForcedBOSTokenLogitsProcessor,
      ForcedEOSTokenLogitsProcessor,
      HammingDiversityLogitsProcessor,
      InfNanRemoveLogitsProcessor,
      LogitsProcessorList,
      MinLengthLogitsProcessor,
      NoBadWordsLogitsProcessor,
      NoRepeatNGramLogitsProcessor,
      PrefixConstrainedLogitsProcessor,
      RepetitionPenaltyLogitsProcessor,
      TemperatureLogitsWarper,
      TopKLogitsWarper,
      TopPLogitsWarper,
  )
  from transformers.generation_stopping_criteria import (
      MaxLengthCriteria,
      MaxTimeCriteria,
      StoppingCriteriaList,
      validate_stopping_criteria,
  )
  from transformers.utils import logging
  from transformers.generation_utils import BeamSearchOutput
  from transformers.generation_utils import BeamSearchEncoderDecoderOutput
  from transformers.generation_utils import BeamSearchDecoderOnlyOutput
  
  
  def beam_search(
          self,
          input_ids: torch.LongTensor,
          beam_scorer: BeamScorer,
          logits_processor: Optional[LogitsProcessorList] = None,
          stopping_criteria: Optional[StoppingCriteriaList] = None,
          max_length: Optional[int] = None,
          pad_token_id: Optional[int] = None,
          eos_token_id: Optional[int] = None,
          output_attentions: Optional[bool] = None,
          output_hidden_states: Optional[bool] = None,
          output_scores: Optional[bool] = None,
          return_dict_in_generate: Optional[bool] = None,
          **model_kwargs,
  ) -> Union[BeamSearchOutput, torch.LongTensor]:
  
      # init values
      logits_processor = logits_processor if logits_processor is not None else LogitsProcessorList()
      stopping_criteria = stopping_criteria if stopping_criteria is not None else StoppingCriteriaList()
      max_length = max_length if max_length is not None else self.config.max_length
      validate_stopping_criteria(stopping_criteria, max_length)
      pad_token_id = pad_token_id if pad_token_id is not None else self.config.pad_token_id
      eos_token_id = eos_token_id if eos_token_id is not None else self.config.eos_token_id
      output_scores = output_scores if output_scores is not None else self.config.output_scores
      output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
      output_hidden_states = (
          output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
      )
      return_dict_in_generate = (
          return_dict_in_generate if return_dict_in_generate is not None else self.config.return_dict_in_generate
      )
  
      # init attention / hidden states / scores tuples
      scores = () if (return_dict_in_generate and output_scores) else None
      decoder_attentions = () if (return_dict_in_generate and output_attentions) else None
      cross_attentions = () if (return_dict_in_generate and output_attentions) else None
      decoder_hidden_states = () if (return_dict_in_generate and output_hidden_states) else None
  
      # if model is an encoder-decoder, retrieve encoder attention weights and hidden states
      if return_dict_in_generate and self.config.is_encoder_decoder:
          encoder_attentions = model_kwargs["encoder_outputs"].get("attentions") if output_attentions else None
          encoder_hidden_states = (
              model_kwargs["encoder_outputs"].get("hidden_states") if output_hidden_states else None
          )
  
      batch_size = len(beam_scorer._beam_hyps)
      num_beams = beam_scorer.num_beams
  
      batch_beam_size, cur_len = input_ids.shape
  
      assert (
              num_beams * batch_size == batch_beam_size
      ), "Batch dimension of `input_ids` should be {num_beams * batch_size}, but is {batch_beam_size}."
  
      beam_scores = torch.zeros((batch_size, num_beams), dtype=torch.float, device=input_ids.device)
      beam_scores[:, 1:] = -1e9
      beam_scores = beam_scores.view((batch_size * num_beams,))
  
      hf_models = model_kwargs.pop("hf_models")
      model_kwargs_list = model_kwargs.pop("model_kwargs_list")
  
      while cur_len < max_length:
          # Changed by JB
          models_inputs = [hf.prepare_inputs_for_generation(input_ids, **m) for hf, m in zip(hf_models, model_kwargs_list)]
          outputs = [hf(**model_inputs,
                        return_dict=True,
                        output_attentions=output_attentions,
                        output_hidden_states=output_hidden_states)
                     for hf, model_inputs in zip(hf_models, models_inputs)]
  
          next_token_logits = sum(o.logits[:, -1, :] for o in outputs)
          ###
  
          # hack: adjust tokens for Marian. For Marian we have to make sure that the `pad_token_id`
          # cannot be generated both before and after the `F.log_softmax` operation.
          next_token_logits = self.adjust_logits_during_generation(
              next_token_logits, cur_len=cur_len, max_length=max_length
          )
  
          next_token_scores = F.log_softmax(next_token_logits, dim=-1)  # (batch_size * num_beams, vocab_size)
  
          next_token_scores = logits_processor(input_ids, next_token_scores)
          next_token_scores = next_token_scores + beam_scores[:, None].expand_as(next_token_scores)
  
          # Store scores, attentions and hidden_states when required
          if return_dict_in_generate:
              if output_scores:
                  scores += (next_token_scores,)
              if output_attentions:
                  decoder_attentions += (
                      (outputs.decoder_attentions,) if self.config.is_encoder_decoder else (outputs.attentions,)
                  )
                  if self.config.is_encoder_decoder:
                      cross_attentions += (outputs.cross_attentions,)
  
              if output_hidden_states:
                  decoder_hidden_states += (
                      (outputs.decoder_hidden_states,)
                      if self.config.is_encoder_decoder
                      else (outputs.hidden_states,)
                  )
  
          # reshape for beam search
          vocab_size = next_token_scores.shape[-1]
          next_token_scores = next_token_scores.view(batch_size, num_beams * vocab_size)
  
          next_token_scores, next_tokens = torch.topk(
              next_token_scores, 2 * num_beams, dim=1, largest=True, sorted=True
          )
  
          next_indices = next_tokens // vocab_size
          next_tokens = next_tokens % vocab_size
  
          # stateless
          beam_outputs = beam_scorer.process(
              input_ids,
              next_token_scores,
              next_tokens,
              next_indices,
              pad_token_id=pad_token_id,
              eos_token_id=eos_token_id,
          )
          beam_scores = beam_outputs["next_beam_scores"]
          beam_next_tokens = beam_outputs["next_beam_tokens"]
          beam_idx = beam_outputs["next_beam_indices"]
  
          input_ids = torch.cat([input_ids[beam_idx, :], beam_next_tokens.unsqueeze(-1)], dim=-1)
  
          cur_len = cur_len + 1
  
          # Changed by JB
          model_kwargs_list = [
              hf._update_model_kwargs_for_generation(o, m, is_encoder_decoder=True)
              for hf, o, m in zip(hf_models, outputs, model_kwargs_list)]
  
          for m in model_kwargs_list:
              if m["past"] is not None:
                  m["past"] = self._reorder_cache(m["past"], beam_idx)
          ###
  
          if beam_scorer.is_done:
              break
  
          if stopping_criteria(input_ids, scores):
              break
  
      sequence_outputs = beam_scorer.finalize(
          input_ids, beam_scores, next_tokens, next_indices, pad_token_id=pad_token_id, eos_token_id=eos_token_id
      )
  
      if return_dict_in_generate:
          if not output_scores:
              sequence_outputs["sequence_scores"] = None
          if self.config.is_encoder_decoder:
              return BeamSearchEncoderDecoderOutput(
                  sequences=sequence_outputs["sequences"],
                  sequences_scores=sequence_outputs["sequence_scores"],
                  scores=scores,
                  encoder_attentions=encoder_attentions,
                  encoder_hidden_states=encoder_hidden_states,
                  decoder_attentions=decoder_attentions,
                  cross_attentions=cross_attentions,
                  decoder_hidden_states=decoder_hidden_states,
              )
          else:
              return BeamSearchDecoderOnlyOutput(
                  sequences=sequence_outputs["sequences"],
                  sequences_scores=sequence_outputs["sequence_scores"],
                  scores=scores,
                  attentions=decoder_attentions,
                  hidden_states=decoder_hidden_states,
              )
      else:
          return sequence_outputs["sequences"]
  