def arc_argmax(parse_probs, length, tokens_to_keep, ensure_tree=True):
if ensure_tree:<EOL><INDENT>I = np.eye(len(tokens_to_keep))<EOL>parse_probs = parse_probs * tokens_to_keep * (<NUM_LIT:1> - I)<EOL>parse_preds = np.argmax(parse_probs, axis=<NUM_LIT:1>)<EOL>tokens = np.arange(<NUM_LIT:1>, length)<EOL>roots = np.where(parse_preds[tokens] == <NUM_LIT:0>)[<NUM_LIT:0>] + <NUM_LIT:1><EOL>if len(roots) < <NUM_LIT:1>:<EOL><INDENT>root_probs = parse_probs[tokens, <NUM_LIT:0>]<EOL>old_head_probs = parse_probs[tokens, parse_preds[tokens]]<EOL>new_root_probs = root_probs / old_head_probs<EOL>new_root = tokens[np.argmax(new_root_probs)]<EOL>parse_preds[new_root] = <NUM_LIT:0><EOL><DEDENT>elif len(roots) > <NUM_LIT:1>:<EOL><INDENT>root_probs = parse_probs[roots, <NUM_LIT:0>]<EOL>parse_probs[roots, <NUM_LIT:0>] = <NUM_LIT:0><EOL>new_heads = np.argmax(parse_probs[roots][:, tokens], axis=<NUM_LIT:1>) + <NUM_LIT:1><EOL>new_head_probs = parse_probs[roots, new_heads] / root_probs<EOL>new_root = roots[np.argmin(new_head_probs)]<EOL>parse_preds[roots] = new_heads<EOL>parse_preds[new_root] = <NUM_LIT:0><EOL><DEDENT>tarjan = Tarjan(parse_preds, tokens)<EOL>for SCC in tarjan.SCCs:<EOL><INDENT>if len(SCC) > <NUM_LIT:1>:<EOL><INDENT>dependents = set()<EOL>to_visit = set(SCC)<EOL>while len(to_visit) > <NUM_LIT:0>:<EOL><INDENT>node = to_visit.pop()<EOL>if not node in dependents:<EOL><INDENT>dependents.add(node)<EOL>to_visit.update(tarjan.edges[node])<EOL><DEDENT><DEDENT>cycle = np.array(list(SCC))<EOL>old_heads = parse_preds[cycle]<EOL>old_head_probs = parse_probs[cycle, old_heads]<EOL>non_heads = np.array(list(dependents))<EOL>parse_probs[np.repeat(cycle, len(non_heads)), np.repeat([non_heads], len(cycle), axis=<NUM_LIT:0>).flatten()] = <NUM_LIT:0><EOL>new_heads = np.argmax(parse_probs[cycle][:, tokens], axis=<NUM_LIT:1>) + <NUM_LIT:1><EOL>new_head_probs = parse_probs[cycle, new_heads] / old_head_probs<EOL>change = np.argmax(new_head_probs)<EOL>changed_cycle = cycle[change]<EOL>old_head = old_heads[change]<EOL>new_head = new_heads[change]<EOL>parse_preds[changed_cycle] = new_head<EOL>tarjan.edges[new_head].add(changed_cycle)<EOL>tarjan.edges[old_head].remove(changed_cycle)<EOL><DEDENT><DEDENT>return parse_preds<EOL><DEDENT>else:<EOL><INDENT>parse_probs = parse_probs * tokens_to_keep<EOL>parse_preds = np.argmax(parse_probs, axis=<NUM_LIT:1>)<EOL>return parse_preds<EOL><DEDENT>
MST
    Adopted from Timothy Dozat https://github.com/tdozat/Parser/blob/master/lib/models/nn.py

    Parameters
    ----------
    parse_probs : NDArray
        seq_len x seq_len, the probability of arcs
    length : NDArray
        real sentence length
    tokens_to_keep : NDArray
        mask matrix
    ensure_tree :
        whether to ensure tree structure of output (apply MST)
    Returns
    -------
    parse_preds : np.ndarray
        prediction of arc parsing with size of (seq_len,)