Fix the relation prediction by heuristic rules

    Parameters
    ----------
    rel_probs : NDArray
        seq_len x rel_size
    length :
        real sentence length
    ensure_tree :
        whether to apply rules
    Returns
    -------
    rel_preds : np.ndarray
        prediction of relations of size (seq_len,)

def rel_argmax(rel_probs, length, ensure_tree=True):
    """Fix the relation prediction by heuristic rules

    Parameters
    ----------
    rel_probs : NDArray
        seq_len x rel_size
    length :
        real sentence length
    ensure_tree :
        whether to apply rules
    Returns
    -------
    rel_preds : np.ndarray
        prediction of relations of size (seq_len,)
    """
    if ensure_tree:
        rel_probs[:, ParserVocabulary.PAD] = 0
        root = ParserVocabulary.ROOT
        tokens = np.arange(1, length)
        rel_preds = np.argmax(rel_probs, axis=1)
        roots = np.where(rel_preds[tokens] == root)[0] + 1
        if len(roots) < 1:
            rel_preds[1 + np.argmax(rel_probs[tokens, root])] = root
        elif len(roots) > 1:
            root_probs = rel_probs[roots, root]
            rel_probs[roots, root] = 0
            new_rel_preds = np.argmax(rel_probs[roots], axis=1)
            new_rel_probs = rel_probs[roots, new_rel_preds] / root_probs
            new_root = roots[np.argmin(new_rel_probs)]
            rel_preds[roots] = new_rel_preds
            rel_preds[new_root] = root
        return rel_preds
    else:
        rel_probs[:, ParserVocabulary.PAD] = 0
        rel_preds = np.argmax(rel_probs, axis=1)
        return rel_preds