def METHOD_NAME(tree_batch, vocab, assm=True):
    set_batch_nodeID(tree_batch, vocab)
    smiles_batch = [tree.smiles for tree in tree_batch]
    jtenc_holder, mess_dict = JTNNEncoder.METHOD_NAME(tree_batch)
    jtenc_holder = jtenc_holder
    mpn_holder = MPN.METHOD_NAME(smiles_batch)
    if assm is False:
        return tree_batch, jtenc_holder, mpn_holder
    cands = []
    batch_idx = []
    for i, mol_tree in enumerate(tree_batch):
        for node in mol_tree.nodes:
            # Leaf node's attachment is determined by neighboring node's attachment
            if node.is_leaf or len(node.cands) == 1:
                continue
            cands.extend([(cand, mol_tree.nodes, node) for cand in node.cands])
            batch_idx.extend([i] * len(node.cands))
    jtmpn_holder = JTMPN.METHOD_NAME(cands, mess_dict)
    batch_idx = torch.LongTensor(batch_idx)
    return tree_batch, jtenc_holder, mpn_holder, (jtmpn_holder, batch_idx)