from __future__ import annotations 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 LabeledMultiSpan, LabeledSpan, MultiSpan, Span from pytorch_ie.core import Document from pytorch_ie.core.document import Annotation, _enumerate_dependencies from src.document.types import ( RelatedRelation, TextDocumentWithLabeledMultiSpansBinaryRelationsLabeledPartitionsAndRelatedRelations, ) from src.utils import distance, distance_slices 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 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