Mask tokens near the end of sequences

[Yux1angJi]

Hi, I am confused about the last part of Forking token selection. It is said you mask tokens near the end of sequences. If some tokens are masked, how can we connect new chains after this node? I also cannot find the implementation of this part in openrlhf/trainer/ppo_utils/entropy_chain_local_manager.py.

# prepare the inference
m_tree_top_n_prompt_ids = []
task_mapping = {}
for i, (tree_idx, node_idx, node, split_idx) in enumerate(expansion_tasks * T):
    prefix_ids = node.get_prefix_ids(split_idx)
    prompt_ids = init_prompt_ids_with_template + prefix_ids
    m_tree_top_n_prompt_ids.append(prompt_ids)
    task_mapping[i] = (tree_idx, node_idx, node, split_idx)

Look forward to your reply, thx.