def mask_unseen_mentions()

in modeling/coval/conll/reader.py [0:0]

• 15 lines of code
• 13 McCabe index (conditional complexity)

def mask_unseen_mentions(clusters, seen_mentions, keep_singletons):
    unseens = {}

    for i, cluster in enumerate(clusters):
        for m in cluster:
            if m not in seen_mentions:
                if i not in unseens:
                    unseens[i] = set()
                unseens[i].add(m)

    remove_clusters = set()
    for i in unseens:
        clusters[i] = [m for m in clusters[i] if m not in unseens[i]]

        if (len(clusters[i]) == 0
                or (len(clusters[i]) == 1 and not keep_singletons)):
            remove_clusters.add(i)

    return [c for i, c in enumerate(clusters) if i not in remove_clusters]