Spaces:
Runtime error
Runtime error
| import logging | |
| from typing import Optional, Sequence, TypeVar, Union | |
| from pie_modules.taskmodules import CrossTextBinaryCorefTaskModule | |
| from pie_modules.taskmodules.cross_text_binary_coref import ( | |
| DocumentType, | |
| SpanDoesNotFitIntoAvailableWindow, | |
| TaskEncodingType, | |
| ) | |
| from pie_modules.utils.tokenization import SpanNotAlignedWithTokenException | |
| from pytorch_ie.annotations import Span | |
| from pytorch_ie.core import TaskEncoding, TaskModule | |
| logger = logging.getLogger(__name__) | |
| S = TypeVar("S", bound=Span) | |
| def shift_span(span: S, offset: int) -> S: | |
| return span.copy(start=span.start + offset, end=span.end + offset) | |
| class CrossTextBinaryCorefTaskModuleWithOptionalContext(CrossTextBinaryCorefTaskModule): | |
| """Same as CrossTextBinaryCorefTaskModule, but: | |
| - optionally without context. | |
| """ | |
| def __init__( | |
| self, | |
| without_context: bool = False, | |
| **kwargs, | |
| ) -> None: | |
| super().__init__(**kwargs) | |
| self.without_context = without_context | |
| def encode_input( | |
| self, | |
| document: DocumentType, | |
| is_training: bool = False, | |
| ) -> Optional[Union[TaskEncodingType, Sequence[TaskEncodingType]]]: | |
| if self.without_context: | |
| return self.encode_input_without_context(document) | |
| else: | |
| return super().encode_input(document) | |
| def encode_input_without_context( | |
| self, document: DocumentType | |
| ) -> Optional[Union[TaskEncodingType, Sequence[TaskEncodingType]]]: | |
| self.collect_all_relations(kind="available", relations=document.binary_coref_relations) | |
| tokenizer_kwargs = dict( | |
| padding=False, | |
| truncation=False, | |
| add_special_tokens=False, | |
| ) | |
| task_encodings = [] | |
| for coref_rel in document.binary_coref_relations: | |
| # TODO: This can miss instances if both texts are the same. We could check that | |
| # coref_rel.head is in document.labeled_spans (same for the tail), but would this | |
| # slow down the encoding? | |
| if not ( | |
| coref_rel.head.target == document.text | |
| or coref_rel.tail.target == document.text_pair | |
| ): | |
| raise ValueError( | |
| f"It is expected that coref relations go from (head) spans over 'text' " | |
| f"to (tail) spans over 'text_pair', but this is not the case for this " | |
| f"relation (i.e. it points into the other direction): {coref_rel.resolve()}" | |
| ) | |
| encoding = self.tokenizer(text=str(coref_rel.head), **tokenizer_kwargs) | |
| encoding_pair = self.tokenizer(text=str(coref_rel.tail), **tokenizer_kwargs) | |
| try: | |
| current_encoding, token_span = self.truncate_encoding_around_span( | |
| encoding=encoding, char_span=shift_span(coref_rel.head, -coref_rel.head.start) | |
| ) | |
| current_encoding_pair, token_span_pair = self.truncate_encoding_around_span( | |
| encoding=encoding_pair, | |
| char_span=shift_span(coref_rel.tail, -coref_rel.tail.start), | |
| ) | |
| except SpanNotAlignedWithTokenException as e: | |
| logger.warning( | |
| f"Could not get token offsets for argument ({e.span}) of coref relation: " | |
| f"{coref_rel.resolve()}. Skip it." | |
| ) | |
| self.collect_relation(kind="skipped_args_not_aligned", relation=coref_rel) | |
| continue | |
| except SpanDoesNotFitIntoAvailableWindow as e: | |
| logger.warning( | |
| f"Argument span [{e.span}] does not fit into available token window " | |
| f"({self.available_window}). Skip it." | |
| ) | |
| self.collect_relation( | |
| kind="skipped_span_does_not_fit_into_window", relation=coref_rel | |
| ) | |
| continue | |
| task_encodings.append( | |
| TaskEncoding( | |
| document=document, | |
| inputs={ | |
| "encoding": current_encoding, | |
| "encoding_pair": current_encoding_pair, | |
| "pooler_start_indices": token_span.start, | |
| "pooler_end_indices": token_span.end, | |
| "pooler_pair_start_indices": token_span_pair.start, | |
| "pooler_pair_end_indices": token_span_pair.end, | |
| }, | |
| metadata={"candidate_annotation": coref_rel}, | |
| ) | |
| ) | |
| self.collect_relation("used", coref_rel) | |
| return task_encodings | |