def expand_word_embedding(tree_old, tree_new, *, vq_codebook_size, num_classes):
  """Add task tokens before current special tokens and after class tokens."""
  cur_embed = tree_new['BertEmbed_0']['word_embeddings']['embedding']
  cur_bias = tree_new['BertMlmLayer_0']['mlm_bias']['bias']
  loaded_embed = tree_old['BertEmbed_0']['word_embeddings']['embedding']
  loaded_bias = tree_old['BertMlmLayer_0']['mlm_bias']['bias']
  assert cur_embed.shape[0] >= loaded_embed.shape[0]
  # Latent and class tokens
  prefix_length = vq_codebook_size + num_classes
  cur_embed = cur_embed.at[:prefix_length].set(loaded_embed[:prefix_length])
  cur_bias = cur_bias.at[:prefix_length].set(loaded_bias[:prefix_length])
  # Other existing special tokens
  suffix_length = loaded_embed.shape[1] - prefix_length
  cur_embed = cur_embed.at[-suffix_length:].set(loaded_embed[-suffix_length:])
  cur_bias = cur_bias.at[-suffix_length:].set(loaded_bias[-suffix_length:])
  tree_old['BertEmbed_0']['word_embeddings']['embedding'] = cur_embed
  tree_old['BertMlmLayer_0']['mlm_bias']['bias'] = cur_bias
  return tree_old
Add task tokens before current special tokens and after class tokens.
videogvt/train_lib/maskgvt_train_lib.py