from __future__ import annotations import itertools import logging from collections import defaultdict from typing import Any, Dict, Iterable, List, Optional, Set, Tuple, TypeVar, Union from pie_modules.utils.span import have_overlap from pytorch_ie import AnnotationLayer from pytorch_ie.annotations import BinaryRelation, LabeledMultiSpan, LabeledSpan, MultiSpan, Span from pytorch_ie.core import Document from pytorch_ie.core.document import Annotation, _enumerate_dependencies from pytorch_ie.documents import TextDocumentWithLabeledSpansAndBinaryRelations from src.document.types import ( RelatedRelation, TextDocumentWithLabeledMultiSpansBinaryRelationsLabeledPartitionsAndRelatedRelations, ) from src.utils import distance, distance_slices from src.utils.graph_utils import get_connected_components from src.utils.span_utils import get_overlap_len logger = logging.getLogger(__name__) D = TypeVar("D", bound=Document) def _remove_overlapping_entities( entities: Iterable[Dict[str, Any]], relations: Iterable[Dict[str, Any]] ) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]: sorted_entities = sorted(entities, key=lambda span: span["start"]) entities_wo_overlap = [] skipped_entities = [] last_end = 0 for entity_dict in sorted_entities: if entity_dict["start"] < last_end: skipped_entities.append(entity_dict) else: entities_wo_overlap.append(entity_dict) last_end = entity_dict["end"] if len(skipped_entities) > 0: logger.warning(f"skipped overlapping entities: {skipped_entities}") valid_entity_ids = set(entity_dict["_id"] for entity_dict in entities_wo_overlap) valid_relations = [ relation_dict for relation_dict in relations if relation_dict["head"] in valid_entity_ids and relation_dict["tail"] in valid_entity_ids ] return entities_wo_overlap, valid_relations def remove_overlapping_entities( doc: D, entity_layer_name: str = "entities", relation_layer_name: str = "relations", ) -> D: # TODO: use document.add_all_annotations_from_other() document_dict = doc.asdict() entities_wo_overlap, valid_relations = _remove_overlapping_entities( entities=document_dict[entity_layer_name]["annotations"], relations=document_dict[relation_layer_name]["annotations"], ) document_dict[entity_layer_name] = { "annotations": entities_wo_overlap, "predictions": [], } document_dict[relation_layer_name] = { "annotations": valid_relations, "predictions": [], } new_doc = type(doc).fromdict(document_dict) return new_doc def remove_partitions_by_labels( document: D, partition_layer: str, label_blacklist: List[str], span_layer: Optional[str] = None ) -> D: """Remove partitions with labels in the blacklist from a document. Args: document: The document to process. partition_layer: The name of the partition layer. label_blacklist: The list of labels to remove. span_layer: The name of the span layer to remove spans from if they are not fully contained in any remaining partition. Any dependent annotations will be removed as well. Returns: The processed document. """ document = document.copy() p_layer: AnnotationLayer = document[partition_layer] new_partitions = [] for partition in p_layer.clear(): if partition.label not in label_blacklist: new_partitions.append(partition) p_layer.extend(new_partitions) if span_layer is not None: result = document.copy(with_annotations=False) removed_span_ids = set() for span in document[span_layer]: # keep spans fully contained in any partition if any( partition.start <= span.start and span.end <= partition.end for partition in new_partitions ): result[span_layer].append(span.copy()) else: removed_span_ids.add(span._id) result.add_all_annotations_from_other( document, removed_annotations={span_layer: removed_span_ids}, strict=False, verbose=False, ) document = result return document D_text = TypeVar("D_text", bound=Document) def remove_annotations_by_label( document: D, layer2label_blacklist: Dict[str, List[str]], verbose: bool = False ) -> D: """Remove annotations with labels in the blacklist from a document. Args: document: The document to process. layer2label_blacklist: A mapping from layer names to lists of labels to remove. verbose: Whether to print number of removed annotations. Returns: The processed document. """ result = document.copy(with_annotations=False) override_annotations: Dict[str, Dict[int, Annotation]] = defaultdict(dict) removed_annotations: Dict[str, Set[int]] = defaultdict(set) for layer_name, labels in layer2label_blacklist.items(): # process gold annotations and predictions for src_layer, tgt_layer in [ (document[layer_name], result[layer_name]), (document[layer_name].predictions, result[layer_name].predictions), ]: current_override_annotations = dict() current_removed_annotations = set() for annotation in src_layer: label = getattr(annotation, "label") if label is None: raise ValueError( f"Annotation {annotation} has no label. Please check the annotation type." ) if label not in labels: current_override_annotations[annotation._id] = annotation.copy() else: current_removed_annotations.add(annotation._id) tgt_layer.extend(current_override_annotations.values()) override_annotations[layer_name].update(current_override_annotations) removed_annotations[layer_name].update(current_removed_annotations) if verbose: num_removed = { layer_name: len(removed_ids) for layer_name, removed_ids in removed_annotations.items() } if len(num_removed) > 0: num_total = { layer_name: len(kept_ids) + num_removed[layer_name] for layer_name, kept_ids in override_annotations.items() } logger.warning( f"doc.id={document.id}: Removed {num_removed} (total: {num_total}) " f"annotations with label blacklists {layer2label_blacklist}" ) result.add_all_annotations_from_other( other=document, removed_annotations=removed_annotations, override_annotations=override_annotations, strict=False, verbose=False, ) return result def replace_substrings_in_text( document: D_text, replacements: Dict[str, str], enforce_same_length: bool = True ) -> D_text: new_text = document.text for old_str, new_str in replacements.items(): if enforce_same_length and len(old_str) != len(new_str): raise ValueError( f'Replacement strings must have the same length, but got "{old_str}" -> "{new_str}"' ) new_text = new_text.replace(old_str, new_str) result_dict = document.asdict() result_dict["text"] = new_text result = type(document).fromdict(result_dict) result.text = new_text return result def replace_substrings_in_text_with_spaces(document: D_text, substrings: Iterable[str]) -> D_text: replacements = {substring: " " * len(substring) for substring in substrings} return replace_substrings_in_text(document, replacements=replacements) def relabel_annotations( document: D, label_mapping: Dict[str, Dict[str, str]], ) -> D: """ Replace annotation labels in a document. Args: document: The document to process. label_mapping: A mapping from layer names to mappings from old labels to new labels. Returns: The processed document. """ dependency_ordered_fields: List[str] = [] _enumerate_dependencies( dependency_ordered_fields, dependency_graph=document._annotation_graph, nodes=document._annotation_graph["_artificial_root"], ) result = document.copy(with_annotations=False) store: Dict[int, Annotation] = {} # not yet used invalid_annotation_ids: Set[int] = set() for field_name in dependency_ordered_fields: if field_name in document._annotation_fields: layer = document[field_name] for is_prediction, anns in [(False, layer), (True, layer.predictions)]: for ann in anns: new_ann = ann.copy_with_store( override_annotation_store=store, invalid_annotation_ids=invalid_annotation_ids, ) if field_name in label_mapping: if ann.label in label_mapping[field_name]: new_label = label_mapping[field_name][ann.label] new_ann = new_ann.copy(label=new_label) else: raise ValueError( f"Label {ann.label} not found in label mapping for {field_name}" ) store[ann._id] = new_ann target_layer = result[field_name] if is_prediction: target_layer.predictions.append(new_ann) else: target_layer.append(new_ann) return result DWithSpans = TypeVar("DWithSpans", bound=Document) def get_start_end(span: Union[Span, MultiSpan]) -> Tuple[int, int]: if isinstance(span, Span): return span.start, span.end elif isinstance(span, MultiSpan): starts, ends = zip(*span.slices) return min(starts), max(ends) else: raise ValueError(f"Unsupported span type: {type(span)}") def _get_aligned_span_mappings( gold_spans: Iterable[Span], pred_spans: Iterable[Span], distance_type: str ) -> Tuple[Dict[int, Span], Dict[int, Span]]: old2new_pred_span = {} span_id2gold_span = {} for pred_span in pred_spans: gold_spans_with_distance = [ ( gold_span, distance( start_end=get_start_end(pred_span), other_start_end=get_start_end(gold_span), distance_type=distance_type, ), ) for gold_span in gold_spans ] if len(gold_spans_with_distance) == 0: continue closest_gold_span, min_distance = min(gold_spans_with_distance, key=lambda x: x[1]) # if the closest gold span is the same as the predicted span, we don't need to align if min_distance == 0.0: continue pred_start_end = get_start_end(pred_span) closest_gold_start_end = get_start_end(closest_gold_span) if have_overlap( start_end=pred_start_end, other_start_end=closest_gold_start_end, ): overlap_len = get_overlap_len(pred_start_end, closest_gold_start_end) l_max = max( pred_start_end[1] - pred_start_end[0], closest_gold_start_end[1] - closest_gold_start_end[0], ) # if the overlap is at least half of the maximum length, we consider it a valid match for alignment valid_match = overlap_len >= (l_max / 2) else: valid_match = False if valid_match: if isinstance(pred_span, Span): aligned_pred_span = pred_span.copy( start=closest_gold_span.start, end=closest_gold_span.end ) elif isinstance(pred_span, MultiSpan): aligned_pred_span = pred_span.copy(slices=closest_gold_span.slices) else: raise ValueError(f"Unsupported span type: {type(pred_span)}") old2new_pred_span[pred_span._id] = aligned_pred_span span_id2gold_span[pred_span._id] = closest_gold_span return old2new_pred_span, span_id2gold_span def get_spans2multi_spans_mapping(multi_spans: Iterable[MultiSpan]) -> Dict[Span, MultiSpan]: result = {} for multi_span in multi_spans: for start, end in multi_span.slices: span_kwargs = dict(start=start, end=end, score=multi_span.score) if isinstance(multi_span, LabeledMultiSpan): result[LabeledSpan(label=multi_span.label, **span_kwargs)] = multi_span else: result[Span(**span_kwargs)] = multi_span return result def align_predicted_span_annotations( document: DWithSpans, span_layer: str, distance_type: str = "center", simple_multi_span: bool = False, verbose: bool = False, ) -> DWithSpans: """ Aligns predicted span annotations with the closest gold spans in a document. First, calculates the distance between each predicted span and each gold span. Then, for each predicted span, the gold span with the smallest distance is selected. If the predicted span and the gold span have an overlap of at least half of the maximum length of the two spans, the predicted span is aligned with the gold span. This also works for MultiSpan annotations, where the slices of the MultiSpan are used to align the predicted spans. If any of the slices is aligned with a gold slice, the MultiSpan is aligned with the respective gold MultiSpan. However, this may result in the predicted MultiSpan being aligned with multiple gold MultiSpans, in which case the closest gold MultiSpan is selected. A simplified version of this alignment can be achieved by setting `simple_multi_span=True`, which treats MultiSpan annotations as simple Spans by using their maximum and minimum start and end indices. Args: document: The document to process. span_layer: The name of the span layer. distance_type: The type of distance to calculate. One of: center, inner, outer simple_multi_span: Whether to treat MultiSpan annotations as simple Spans by using their maximum and minimum start and end indices. verbose: Whether to print debug information. Returns: The processed document. """ gold_spans = document[span_layer] if len(gold_spans) == 0: return document.copy() pred_spans = document[span_layer].predictions span_annotation_type = document.annotation_types()[span_layer] if issubclass(span_annotation_type, Span) or simple_multi_span: old2new_pred_span, span_id2gold_span = _get_aligned_span_mappings( gold_spans=gold_spans, pred_spans=pred_spans, distance_type=distance_type ) elif issubclass(span_annotation_type, MultiSpan): # create Span objects from MultiSpan slices gold_single_spans2multi_spans = get_spans2multi_spans_mapping(gold_spans) pred_single_spans2multi_spans = get_spans2multi_spans_mapping(pred_spans) # create the alignment mappings for the single spans single_old2new_pred_span, single_span_id2gold_span = _get_aligned_span_mappings( gold_spans=gold_single_spans2multi_spans.keys(), pred_spans=pred_single_spans2multi_spans.keys(), distance_type=distance_type, ) # collect all Spans that are part of the same MultiSpan pred_multi_span2single_spans: Dict[MultiSpan, List[Span]] = defaultdict(list) for pred_span, multi_span in pred_single_spans2multi_spans.items(): pred_multi_span2single_spans[multi_span].append(pred_span) # create the new mappings for the MultiSpans old2new_pred_span = {} span_id2gold_span = {} for pred_multi_span, pred_single_spans in pred_multi_span2single_spans.items(): # if any of the single spans is aligned with a gold span, align the multi span if any( pred_single_span._id in single_old2new_pred_span for pred_single_span in pred_single_spans ): # get aligned gold multi spans aligned_gold_multi_spans = set() for pred_single_span in pred_single_spans: if pred_single_span._id in single_old2new_pred_span: aligned_gold_single_span = single_span_id2gold_span[pred_single_span._id] aligned_gold_multi_span = gold_single_spans2multi_spans[ aligned_gold_single_span ] aligned_gold_multi_spans.add(aligned_gold_multi_span) # calculate distances between the predicted multi span and the aligned gold multi spans gold_multi_spans_with_distance = [ ( gold_multi_span, distance_slices( slices=pred_multi_span.slices, other_slices=gold_multi_span.slices, distance_type=distance_type, ), ) for gold_multi_span in aligned_gold_multi_spans ] if len(aligned_gold_multi_spans) > 1: logger.warning( f"Multiple gold multi spans aligned with predicted multi span ({pred_multi_span}): " f"{aligned_gold_multi_spans}" ) # get the closest gold multi span closest_gold_multi_span, min_distance = min( gold_multi_spans_with_distance, key=lambda x: x[1] ) old2new_pred_span[pred_multi_span._id] = pred_multi_span.copy( slices=closest_gold_multi_span.slices ) span_id2gold_span[pred_multi_span._id] = closest_gold_multi_span else: raise ValueError(f"Unsupported span annotation type: {span_annotation_type}") result = document.copy(with_annotations=False) # multiple predicted spans can be aligned with the same gold span, # so we need to keep track of the added spans added_pred_span_ids = dict() for pred_span in pred_spans: # just add the predicted span if it was not aligned with a gold span if pred_span._id not in old2new_pred_span: # if this was not added before (e.g. as aligned span), add it if pred_span._id not in added_pred_span_ids: keep_pred_span = pred_span.copy() result[span_layer].predictions.append(keep_pred_span) added_pred_span_ids[pred_span._id] = keep_pred_span elif verbose: print(f"Skipping duplicate predicted span. pred_span='{str(pred_span)}'") else: aligned_pred_span = old2new_pred_span[pred_span._id] # if this was not added before (e.g. as aligned or original pred span), add it if aligned_pred_span._id not in added_pred_span_ids: result[span_layer].predictions.append(aligned_pred_span) added_pred_span_ids[aligned_pred_span._id] = aligned_pred_span elif verbose: prev_pred_span = added_pred_span_ids[aligned_pred_span._id] gold_span = span_id2gold_span[pred_span._id] print( f"Skipping duplicate aligned predicted span. aligned gold_span='{str(gold_span)}', " f"prev_pred_span='{str(prev_pred_span)}', current_pred_span='{str(pred_span)}'" ) # print("bbb") result[span_layer].extend([span.copy() for span in gold_spans]) # add remaining gold and predicted spans (the result, _aligned_spans, is just for debugging) _aligned_spans = result.add_all_annotations_from_other( document, override_annotations={span_layer: old2new_pred_span} ) return result def add_related_relations_from_binary_relations( document: TextDocumentWithLabeledMultiSpansBinaryRelationsLabeledPartitionsAndRelatedRelations, link_relation_label: str, link_partition_whitelist: Optional[List[List[str]]] = None, relation_label_whitelist: Optional[List[str]] = None, reversed_relation_suffix: str = "_reversed", symmetric_relations: Optional[List[str]] = None, ) -> TextDocumentWithLabeledMultiSpansBinaryRelationsLabeledPartitionsAndRelatedRelations: span2partition = {} for multi_span in document.labeled_multi_spans: found_partition = False for partition in document.labeled_partitions or [ LabeledSpan(start=0, end=len(document.text), label="ALL") ]: starts, ends = zip(*multi_span.slices) if partition.start <= min(starts) and max(ends) <= partition.end: span2partition[multi_span] = partition found_partition = True break if not found_partition: raise ValueError(f"No partition found for multi_span {multi_span}") rel_head2rels = defaultdict(list) rel_tail2rels = defaultdict(list) for rel in document.binary_relations: rel_head2rels[rel.head].append(rel) rel_tail2rels[rel.tail].append(rel) link_partition_whitelist_tuples = None if link_partition_whitelist is not None: link_partition_whitelist_tuples = {tuple(pair) for pair in link_partition_whitelist} skipped_labels = [] for link_rel in document.binary_relations: if link_rel.label == link_relation_label: head_partition = span2partition[link_rel.head] tail_partition = span2partition[link_rel.tail] if link_partition_whitelist_tuples is None or ( (head_partition.label, tail_partition.label) in link_partition_whitelist_tuples ): # link_head -> link_tail == rel_head -> rel_tail for rel in rel_head2rels.get(link_rel.tail, []): label = rel.label if relation_label_whitelist is None or label in relation_label_whitelist: new_rel = RelatedRelation( head=link_rel.head, tail=rel.tail, link_relation=link_rel, relation=rel, label=label, ) document.related_relations.append(new_rel) else: skipped_labels.append(label) # link_head -> link_tail == rel_tail -> rel_head if reversed_relation_suffix is not None: for reversed_rel in rel_tail2rels.get(link_rel.tail, []): label = reversed_rel.label if not (symmetric_relations is not None and label in symmetric_relations): label = f"{label}{reversed_relation_suffix}" if relation_label_whitelist is None or label in relation_label_whitelist: new_rel = RelatedRelation( head=link_rel.head, tail=reversed_rel.head, link_relation=link_rel, relation=reversed_rel, label=label, ) document.related_relations.append(new_rel) else: skipped_labels.append(label) else: logger.warning( f"Skipping related relation because of partition whitelist ({[head_partition.label, tail_partition.label]}): {link_rel.resolve()}" ) if len(skipped_labels) > 0: logger.warning( f"Skipped relations with labels not in whitelist: {sorted(set(skipped_labels))}" ) return document T = TypeVar("T", bound=TextDocumentWithLabeledSpansAndBinaryRelations) def remove_discontinuous_spans( document: T, parts_of_same_relation: str, verbose: bool = False, ) -> T: """ Remove discontinuous spans from a document. Args: document: The document to process. parts_of_same_relation: The name of the relation that indicates linked spans. verbose: Whether to print debug information. Returns: The processed document. """ result = document.copy() spans = result.labeled_spans.clear() rels = result.binary_relations.clear() segment_spans = set() segment_rels = set() # collect all spans that are linked for rel in rels: if rel.label == parts_of_same_relation: segment_spans.add(rel.head) segment_spans.add(rel.tail) segment_rels.add(rel) for span in spans: if span not in segment_spans: result.labeled_spans.append(span) other_rels_dropped = set() for rel in rels: if rel not in segment_rels: if rel.head not in segment_spans and rel.tail not in segment_spans: result.binary_relations.append(rel) else: other_rels_dropped.add(rel) if verbose: if len(segment_rels) > 0: logger.warning( f"doc={document.id}: Dropped {len(segment_rels)} segment rels " f"and {len(other_rels_dropped)} other rels " f"({round((len(document.binary_relations) - len(result.binary_relations)) * 100 / len(document.binary_relations), 1)}% " f"of all relations dropped)" ) return result def close_clusters_transitively( document: D, relation_layer: str, link_relation_label: str, verbose: bool = False ) -> D: """ Close clusters transitively by adding relations between all pairs of spans in the same cluster. Args: document: The document to process. relation_layer: The name of the relation layer. link_relation_label: The label of the link relation. verbose: Whether to print debug information. Returns: The processed document. """ result = document.copy() connected_components: List[List[Annotation]] = get_connected_components( relations=result[relation_layer], link_relation_label=link_relation_label, add_singletons=False, ) # detach from document relations = result[relation_layer].clear() # use set to speed up membership checks relations_set = set(relations) n_before = len(relations) for cluster in connected_components: for head, tail in itertools.combinations(sorted(cluster), 2): rel = BinaryRelation( head=head, tail=tail, label=link_relation_label, ) rel_reversed = BinaryRelation( head=tail, tail=head, label=link_relation_label, ) if rel not in relations_set and rel_reversed not in relations_set: # append to relations to keep the order relations.append(rel) relations_set.add(rel) result[relation_layer].extend(relations) if verbose: num_added = len(relations) - n_before if num_added > 0: logger.warning( f"doc.id={document.id}: added {num_added} relations to {relation_layer} layer" ) return result def get_ancestor_layers(children: Dict[str, Set[str]], layer: str) -> Set[str]: """ Get all ancestor layers of a given layer in the dependency graph. Args: children: A mapping from layers to their children layers. layer: The layer for which to find ancestors. Returns: A set of ancestor layers. """ ancestors = set() def _get_ancestors(current_layer: str): for parent_layer, child_layers in children.items(): if current_layer in child_layers: ancestors.add(parent_layer) _get_ancestors(parent_layer) _get_ancestors(layer) # drop the _artificial_root ancestors.discard("_artificial_root") return ancestors def remove_binary_relations_by_partition_labels( document: D, partition_layer: str, relation_layer: str, partition_label_whitelist: Optional[List[List[str]]] = None, partition_label_blacklist: Optional[List[List[str]]] = None, verbose: bool = False, ) -> D: """ Remove binary relations that are not between partitions with labels in the whitelist or that are in the blacklist. Args: document: The document to process. partition_layer: The name of the partition layer. relation_layer: The name of the relation layer. partition_label_whitelist: The list of head-tail label pairs to keep. partition_label_blacklist: The list of head-tail label pairs to remove. verbose: Whether to print the removed relations to console. Returns: The processed document. """ result = document.copy() relation_annotation_layer = result[relation_layer] # get all layers that target the relation layer relation_dependent_layers = get_ancestor_layers( children=result._annotation_graph, layer=relation_layer ) # clear all layers that depend on the relation layer for layer_name in relation_dependent_layers: dependent_layer = result[layer_name] gold_anns_cleared = dependent_layer.clear() pred_anns_cleared = dependent_layer.predictions.clear() if len(gold_anns_cleared) > 0 or len(pred_anns_cleared) > 0: if verbose: logger.warning( f"doc.id={document.id}: Cleared {len(gold_anns_cleared)} gold and " f"{len(pred_anns_cleared)} predicted annotations from layer {layer_name} " f"because it depends on the relation layer {relation_layer}." ) span2partition = {} span_layer: AnnotationLayer for span_layer in relation_annotation_layer.target_layers.values(): for span in list(span_layer) + list(span_layer.predictions): if isinstance(span, Span): span_start, span_end = span.start, span.end elif isinstance(span, MultiSpan): span_start, span_end = min(start for start, _ in span.slices), max( end for _, end in span.slices ) else: raise ValueError(f"Unsupported span type: {type(span)}") found_partition = False for partition in result[partition_layer]: if partition.start <= span_start and span_end <= partition.end: span2partition[span] = partition found_partition = True break if not found_partition: raise ValueError(f"No partition found for span {span}") if partition_label_whitelist is not None: partition_label_whitelist_tuples = [tuple(pair) for pair in partition_label_whitelist] else: partition_label_whitelist_tuples = None if partition_label_blacklist is not None: partition_label_blacklist_tuples = [tuple(pair) for pair in partition_label_blacklist] else: partition_label_blacklist_tuples = None for relation_base_layer in [relation_annotation_layer, relation_annotation_layer.predictions]: # get all relations and clear the layer relations = relation_base_layer.clear() for relation in relations: head_partition = span2partition[relation.head] tail_partition = span2partition[relation.tail] pair = (head_partition.label, tail_partition.label) if ( partition_label_whitelist_tuples is None or pair in partition_label_whitelist_tuples ) and ( partition_label_blacklist_tuples is None or pair not in partition_label_blacklist_tuples ): relation_base_layer.append(relation) else: if verbose: logger.info( f"Removing relation {relation} because its partitions " f"({pair}) are not in the whitelist or are in the blacklist." ) return result