docstr: # adapted from MegatronGPTModel._bucketize_gpt_inference 
docstr_tok: ["#","adapted","from","MegatronGPTModel._bucketize_gpt_inference"]
code: def prepare_megatron_generation(
      self, labels, input_ids, template_length
  ):
      """
      # adapted from MegatronGPTModel._bucketize_gpt_inference
      """
      batch_size = labels.size(0)
      prompt_tags = (
          [self.prompt_tags[0]] * batch_size
          if self.prompt_tags
          else None
      )
      batch_tokens = input_ids.tolist()
  
      # unpad tokens
      lens = template_length
      indxs = [index for index in range(batch_size)]
      for lenn, index in zip(lens, indxs):
          batch_tokens[index] = batch_tokens[index][:lenn]
  
      # chunk tokens by same length
      pre_buckets, lens = [], list(set(lens.tolist()))
      for lenn in lens:
          pre_buckets.append(
              [
                  (tokens, index)
                  for index, tokens in enumerate(batch_tokens)
                  if len(tokens) == lenn
              ]
          )
  
      buckets, positions, bucket_prompt_tags = [], [], []
  
      # get buckets and prompts initial positions
      for bucket in pre_buckets:
          buckets.append(
              torch.tensor([item[0] for item in bucket]).to(
                  device=labels.device
              )
          )
          positions.append([item[1] for item in bucket])
  
          # bucket prompt tags identically to their corresponding examples
          if prompt_tags:
              bucket_prompt_tags.append(
                  [prompt_tags[item[1]] for item in bucket]
              )
  
      # Flatten position list
      positions = [item for sublist in positions for item in sublist]
  
      # Flatten buckets and bucket_prompt_tags # temp fix for megatron complete issue. However, this is also slower than bucketized inference
      buckets = [
          item.unsqueeze(0) for sublist in buckets for item in sublist
      ]
      bucket_prompt_tags = [
          [item] for sublist in bucket_prompt_tags for item in sublist
      ]
  
      request = {"tokens": buckets, "prompt_tags": bucket_prompt_tags}
  
      return positions, request
  
imports: import os
  from typing import Dict, Optional, Union
  import numpy
  import torch
  from omegaconf import DictConfig, OmegaConf, open_dict
  from pytorch_lightning import Trainer
  from torch.utils.data import DataLoader
  from transformers import AutoModelForSeq2SeqLM
  from nemo.collections.nlp.data.dialogue import DialogueSGDDataProcessor
  from nemo.collections.nlp.data.dialogue.data_processor.mellon_qa_data_processor import DialogueMellonQADataProcessor
  from nemo.collections.nlp.data.dialogue.data_processor.ms_marco_data_processor import DialogueMSMarcoDataProcessor
  from nemo.collections.nlp.data.dialogue.dataset.dialogue_s2s_generation_dataset import DialogueS2SGenerationDataset
  from nemo.collections.nlp.metrics.dialogue_metrics import DialogueGenerationMetrics
  from nemo.collections.nlp.models.language_modeling.megatron_t5_model import MegatronT5Model
  from nemo.collections.nlp.models.nlp_model import NLPModel
  from nemo.core.classes.common import PretrainedModelInfo
  from nemo.utils import logging
label_desc: Use image node_id 10 for calling the DialogueS2SGenerationModel obj's underlying member method code with example usage: obj.prepare_megatron_generation(labels, input_ids, template_length) and returns: positions,  request