def generate_tokens_for_lm(batch: Batch,
                           task: task_registry.Task,
                           config: ml_collections.ConfigDict,):
  """Generate input tokens for language model.

  Args:
    batch: A single batch of data, where batch['inputs'] is latent code of shape
      [device_bs, l_t, l_h, l_w] in jnp.int32, batch['label'] is class id of
      shape [device_bs, 1] in jnp.int32 when config.class_conditional is True,
      and batch['cond_inputs'] is latent code like batch['inputs'] when
      config.frame_conditional is True.
    task: current task.
    config: model config.

  Returns:
    Dictionary of the masked and target token sequence with the class info.
  """
  batch_size = batch['inputs'].shape[0]
  tokens_list = []

  assert not config.get('multi_task', False)
  # Task condition token
  if config.get('multi_task', False):
    task_tokens = jnp.full((batch_size, 1), task.task_token, jnp.int32)
    tokens_list.append(mask_utils.no_mask({'inputs': task_tokens}))

  # Class condition token
  class_tokens = get_class_tokens(batch, config, task.config)
  tokens_list.append(mask_utils.no_mask({'inputs': class_tokens}))

  # Frame condition tokens
  if (task.task_id == 'frame_prediction' and task.config.get(
      'condition_mode') == 'prefix'):
    num_frames = config.frame_prediction.cond_latent_frames
    cond_latents = batch.pop('cond_inputs')[:, :num_frames]
    tokens_list.append(mask_utils.no_mask({'inputs': cond_latents}))
    sep_token_ids = jnp.full((batch_size, 1), task_registry.SpecialTokens.SEP,
                             jnp.int32)
    tokens_list.append(mask_utils.no_mask({'inputs': sep_token_ids}))
    batch.pop('cond_mask')

  # Latent tokens
  latent_tokens = mask_utils.no_mask(batch)
  tokens_list.append(latent_tokens)
  all_tokens = mask_utils.concat_tokens(*tokens_list)
  # no mask tokens in lm, so here use mask token as bos token.
  all_tokens = mask_utils.lm_tokens(all_tokens['masked_inputs'],
                                    task_registry.SpecialTokens.MASK)
  cond_tokens = mask_utils.concat_tokens(*(tokens_list[:-1]))
  # 1 for bos
  all_tokens['cond_lengths'] = cond_tokens['masked_inputs'].shape[-1] + 1
  return all_tokens
Generate input tokens for language model.

  Args:
    batch: A single batch of data, where batch['inputs'] is latent code of shape
      [device_bs, l_t, l_h, l_w] in jnp.int32, batch['label'] is class id of
      shape [device_bs, 1] in jnp.int32 when config.class_conditional is True,
      and batch['cond_inputs'] is latent code like batch['inputs'] when
      config.frame_conditional is True.
    task: current task.
    config: model config.

  Returns:
    Dictionary of the masked and target token sequence with the class info.
  
videogvt/train_lib/maskgvt_train_lib.py