| """KEEP/MERGE pair 출력의 deterministic transitive clustering.""" |
|
|
| from __future__ import annotations |
|
|
| from collections import defaultdict |
| from typing import Iterable |
|
|
|
|
| class UnionFindClusterer: |
| def cluster( |
| self, |
| mention_ids: Iterable[str], |
| merge_edges: Iterable[tuple[str, str, float]], |
| *, |
| threshold: float = 0.5, |
| ) -> tuple[tuple[str, ...], ...]: |
| parent = {item: item for item in mention_ids} |
|
|
| def find(item: str) -> str: |
| while parent[item] != item: |
| parent[item] = parent[parent[item]] |
| item = parent[item] |
| return item |
|
|
| for left, right, score in merge_edges: |
| if score < threshold: |
| continue |
| if left not in parent or right not in parent: |
| raise ValueError("coreference edge references unknown mention") |
| left_root, right_root = find(left), find(right) |
| if left_root != right_root: |
| parent[right_root] = left_root |
| groups: dict[str, list[str]] = defaultdict(list) |
| for item in parent: |
| groups[find(item)].append(item) |
| return tuple(sorted((tuple(sorted(group)) for group in groups.values()), key=lambda group: group[0])) |
|
|
