-
Notifications
You must be signed in to change notification settings - Fork 1
/
tokenization.py
515 lines (461 loc) · 22.9 KB
/
tokenization.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
import functools
import json
import logging
from collections import defaultdict
from copy import copy, deepcopy
from typing import (
Callable,
Dict,
Iterable,
List,
Optional,
Set,
Tuple,
Type,
TypeVar,
Union,
)
from pytorch_ie.core import Annotation
from pytorch_ie.documents import TextBasedDocument, TokenBasedDocument
from transformers import PreTrainedTokenizer
from pie_modules.annotations import MultiSpan, Span
from pie_modules.utils import resolve_type
logger = logging.getLogger(__name__)
ToD = TypeVar("ToD", bound=TokenBasedDocument)
TeD = TypeVar("TeD", bound=TextBasedDocument)
def find_token_offset_mapping(text: str, tokens: Iterable[str]) -> List[Tuple[int, int]]:
"""Find the token offset mapping for a given text and tokens. If a token is not found in the
text, the token offset mapping will be (idx, idx) for this token. So, this works also if
special tokens are not part of the text.
Args:
text (str): The text.
tokens (Iterable[str]): The tokens.
Returns:
List[Tuple[int, int]]: The token offset mapping.
"""
token_offset_mapping = []
start = 0
for token in tokens:
new_start = text.find(token, start)
if new_start == -1:
token_offset_mapping.append((start, start))
continue
end = new_start + len(token)
token_offset_mapping.append((new_start, end))
start = end
return token_offset_mapping
def char_span_to_token_span(
span: Annotation, char_to_token: Callable[[int], Optional[int]]
) -> Optional[Union[Span, MultiSpan]]:
if isinstance(span, Span):
start_token_idx = char_to_token(span.start)
end_token_idx_inclusive = char_to_token(span.end - 1)
if start_token_idx is None or end_token_idx_inclusive is None:
return None
return span.copy(start=start_token_idx, end=end_token_idx_inclusive + 1)
elif isinstance(span, MultiSpan):
slices_inclusive_end = [
(char_to_token(start), char_to_token(end - 1)) for start, end in span.slices
]
if any(start is None or end is None for start, end in slices_inclusive_end):
return None
return span.copy(
slices=tuple(
# ignore type because we checked that start and end are not None
(start, inclusive_end + 1) # type: ignore
for start, inclusive_end in slices_inclusive_end
)
)
else:
raise TypeError(
f"can not convert layers that target the text but contain non-span annotations, but found {type(span)}"
)
def token_span_to_char_span(
span: Annotation, token_offset_mapping: List[Tuple[int, int]]
) -> Optional[Union[Span, MultiSpan]]:
if isinstance(span, Span):
start_char_idx = token_offset_mapping[span.start][0]
end_char_idx = token_offset_mapping[span.end - 1][1]
return span.copy(start=start_char_idx, end=end_char_idx)
elif isinstance(span, MultiSpan):
slices = [
(token_offset_mapping[start][0], token_offset_mapping[end - 1][1])
for start, end in span.slices
]
return span.copy(slices=slices)
else:
raise TypeError(
f"can not convert layers that target the tokens but contain non-span annotations, but found {type(span)}"
)
def span_sort_key(span: Union[Span, MultiSpan]) -> Tuple[int, ...]:
if isinstance(span, Span):
return span.start, span.end
elif isinstance(span, MultiSpan):
result: List[int] = []
for start, end in span.slices:
result.extend((start, end))
return tuple(result)
else:
raise TypeError(f"can not sort {type(span)}")
def text_based_document_to_token_based(
doc: TextBasedDocument,
result_document_type: Union[Type[ToD], str],
tokens: Optional[List[str]] = None,
token_offset_mapping: Optional[List[Tuple[int, int]]] = None,
char_to_token: Optional[Callable[[int], Optional[int]]] = None,
strict_span_conversion: bool = True,
verbose: bool = True,
added_annotations: Optional[Dict[str, Dict[Annotation, Annotation]]] = None,
) -> ToD:
"""Convert a text based document to a token based document. Uses either tokens,
token_offset_mapping or char_to_token (provided explicitly or from the document metadata) to
convert the annotations that target the text and all that depend on these (also adds all
remaining annotations).
Args:
doc (TextBasedDocument): The text based document.
result_document_type (Union[Type[ToD], str]): The type of the token based document.
tokens (Optional[List[str]], optional): The tokens. If None, the tokens are taken from the
metadata. Defaults to None.
token_offset_mapping (Optional[List[Tuple[int, int]]], optional): The token offset mapping.
If None, the token offset mapping is taken from the metadata. Defaults to None.
char_to_token (Optional[Callable[[int], Optional[int]]], optional): The char to token function.
If None, the char to token function is constructed from the token offset mapping. Defaults
to None.
strict_span_conversion (bool, optional): If True, raise an error if not all annotations can
be converted to token based documents. Defaults to True.
verbose (bool, optional): If True, log warnings if annotations can not be converted. Defaults
to True.
added_annotations (Optional[Dict[str, Dict[Annotation, Annotation]]], optional): Pass an empty
dictionary to collect the added annotations. Defaults to None.
Returns:
ToD: The token based document of type result_document_type with the converted annotations.
"""
document_type = resolve_type(
type_or_str=result_document_type, expected_super_type=TokenBasedDocument
)
metadata = deepcopy(doc.metadata)
if tokens is None:
tokens = doc.metadata.get("tokens")
elif "tokens" in metadata and metadata["tokens"] != tokens:
logger.warning("tokens in metadata are different from new tokens, take the new tokens")
# save text, token_offset_mapping and char_to_token (if available) in metadata
metadata["text"] = doc.text
token_offset_mapping_lists: Optional[List[List[int]]]
if token_offset_mapping is None:
token_offset_mapping_lists = metadata.get("token_offset_mapping")
if token_offset_mapping_lists is not None:
token_offset_mapping = [tuple(offsets) for offsets in token_offset_mapping_lists] # type: ignore
else:
# convert offset tuples to lists because serialization and deserialization again
# will produce lists in any way (json does not know tuples)
token_offset_mapping_lists = [list(offsets) for offsets in token_offset_mapping]
if (
"token_offset_mapping" in metadata
and metadata["token_offset_mapping"] != token_offset_mapping_lists
):
logger.warning(
"token_offset_mapping in metadata is different from the new token_offset_mapping, "
"overwrite the metadata"
)
metadata["token_offset_mapping"] = token_offset_mapping_lists
if tokens is None:
if token_offset_mapping is not None:
tokens = [doc.text[start:end] for start, end in token_offset_mapping]
else:
raise ValueError(
"tokens or token_offset_mapping must be provided to convert a text based document to token based, "
"but got None for both"
)
if char_to_token is None:
char_to_token = metadata.get("char_to_token")
else:
if "char_to_token" in metadata and metadata["char_to_token"] != char_to_token:
logger.warning(
"char_to_token in metadata is different from the new char_to_token, overwrite the metadata"
)
metadata["char_to_token"] = char_to_token
# construct the char_to_token function, if not provided, from the token_offset_mapping
if char_to_token is None:
if token_offset_mapping is None:
token_offset_mapping = find_token_offset_mapping(text=doc.text, tokens=tokens)
char_to_token_dict: Dict[int, int] = {}
for token_idx, (start, end) in enumerate(token_offset_mapping):
for char_idx in range(start, end):
char_to_token_dict[char_idx] = token_idx
def char_to_token(char_idx: int) -> Optional[int]:
return char_to_token_dict.get(char_idx)
result = document_type(tokens=tuple(tokens), id=doc.id, metadata=metadata)
text_targeting_layers = [
annotation_field.name
for annotation_field in doc.annotation_fields()
if "text" in annotation_field.metadata["targets"]
]
override_annotations: Dict[str, Dict[int, Annotation]] = {}
removed_annotations: Dict[str, Set[int]] = defaultdict(set)
for text_targeting_layer_name in text_targeting_layers:
override_annotations[text_targeting_layer_name] = {}
for char_span in doc[text_targeting_layer_name]:
token_span = char_span_to_token_span(char_span, char_to_token)
if token_span is None:
if strict_span_conversion:
raise ValueError(
f'cannot find token span for character span: "{char_span}", text="{doc.text}", '
f"token_offset_mapping={token_offset_mapping}"
)
else:
if verbose:
logger.warning(
f'cannot find token span for character span "{char_span}", skip it (disable this '
f"warning with verbose=False)"
)
removed_annotations[text_targeting_layer_name].add(char_span._id)
else:
override_annotations[text_targeting_layer_name][char_span._id] = token_span
if added_annotations is not None:
added_annotations.setdefault(text_targeting_layer_name, {})[
char_span
] = token_span
valid_spans = set(override_annotations[text_targeting_layer_name].values())
result[text_targeting_layer_name].extend(sorted(valid_spans, key=span_sort_key))
added_annotations_from_remaining_layers = result.add_all_annotations_from_other(
doc,
override_annotations=override_annotations,
removed_annotations=removed_annotations,
strict=strict_span_conversion,
verbose=verbose,
)
if added_annotations is not None:
for layer_name, orig_ann_id2new_ann in added_annotations_from_remaining_layers.items():
ann_id2ann = {
ann._id: ann for ann in list(doc[layer_name]) + list(doc[layer_name].predictions)
}
annotation_mapping = {
ann_id2ann[orig_ann_id]: new_ann
for orig_ann_id, new_ann in orig_ann_id2new_ann.items()
}
added_annotations.setdefault(layer_name, {}).update(annotation_mapping)
return result
def token_based_document_to_text_based(
doc: TokenBasedDocument,
result_document_type: Union[Type[TeD], str],
text: Optional[str] = None,
token_offset_mapping: Optional[List[Tuple[int, int]]] = None,
join_tokens_with: Optional[str] = None,
strict_span_conversion: bool = True,
verbose: bool = True,
added_annotations: Optional[Dict[str, Dict[Annotation, Annotation]]] = None,
) -> TeD:
"""Convert a token based document to a text based document. Uses either text,
token_offset_mapping or char_to_token (provided explicitly or from the document metadata) to
convert the annotations that target the tokens and all that depend on these (also adds all
remaining annotations).
Args:
doc (TokenBasedDocument): The token based document.
result_document_type (Union[Type[TeD], str]): The type of the text based document.
text (Optional[str], optional): The text. If None, constructed from the tokens (requires
join_tokens_with) or taken from the metadata. Defaults to None.
token_offset_mapping (Optional[List[Tuple[int, int]]], optional): The token offset mapping.
If None, the token offset mapping is constructed from the tokens (requires join_tokens_with)
or taken from the metadata. Defaults to None.
join_tokens_with (Optional[str], optional): The token separator. If no text is provided, the
text and token offset mapping are constructed from the tokens by joining them with this
separator. Defaults to None.
strict_span_conversion (bool, optional): If True, raise an error if not all annotations
can be converted to text based documents. Defaults to True.
verbose (bool, optional): If True, log warnings if annotations can not be converted.
Defaults to True.
added_annotations (Optional[Dict[str, Dict[Annotation, Annotation]]], optional): Pass an
empty dictionary to collect the added annotations. Defaults to None.
Returns:
TeD: The text based document of type result_document_type with the converted annotations.
"""
document_type = resolve_type(
type_or_str=result_document_type, expected_super_type=TextBasedDocument
)
# if a token_separator is provided, we construct the text from the tokens
if text is None and join_tokens_with is not None:
start = 0
token_offset_mapping = []
tokens = doc.tokens
for token in tokens:
end = start + len(token)
token_offset_mapping.append((start, end))
# we add the separator after each token
start = end + len(join_tokens_with)
text = join_tokens_with.join(tokens)
# otherwise we try to use the text from the metadata
if text is None:
text = doc.metadata.get("text")
if text is None:
raise ValueError(
"if join_tokens_with is None, text must be provided, but got None as well"
)
token_offset_mapping_lists = (
doc.metadata.get("token_offset_mapping")
if token_offset_mapping is None
else token_offset_mapping
)
if token_offset_mapping_lists is None:
token_offset_mapping = find_token_offset_mapping(text=text, tokens=doc.tokens)
else:
# we convert the token_offset_mapping to tuples because the token_offset_mapping
# in the metadata is a list of lists, but we need a list of tuples
token_offset_mapping = [tuple(offsets) for offsets in token_offset_mapping_lists] # type: ignore
result = document_type(text=text, id=doc.id, metadata=deepcopy(doc.metadata))
result.metadata["tokens"] = list(doc.tokens)
# convert offset tuples to lists because serialization and deserialization again
# will produce lists in any way (json does not know tuples)
token_offset_mapping_lists = [list(offsets) for offsets in token_offset_mapping]
if (
"token_offset_mapping" in doc.metadata
and doc.metadata["token_offset_mapping"] != token_offset_mapping_lists
):
logger.warning(
"token_offset_mapping in metadata is different from the new token_offset_mapping, "
"overwrite the metadata"
)
result.metadata["token_offset_mapping"] = token_offset_mapping_lists
token_targeting_layers = [
annotation_field.name
for annotation_field in doc.annotation_fields()
if "tokens" in annotation_field.metadata["targets"]
]
override_annotations: Dict[str, Dict[int, Annotation]] = {}
removed_annotations: Dict[str, Set[int]] = defaultdict(set)
for token_targeting_layer_name in token_targeting_layers:
override_annotations[token_targeting_layer_name] = {}
for token_span in doc[token_targeting_layer_name]:
char_span = token_span_to_char_span(token_span, token_offset_mapping)
override_annotations[token_targeting_layer_name][token_span._id] = char_span
if added_annotations is not None:
added_annotations.setdefault(token_targeting_layer_name, {})[
token_span
] = char_span
valid_spans = set(override_annotations[token_targeting_layer_name].values())
result[token_targeting_layer_name].extend(sorted(valid_spans, key=span_sort_key))
added_annotations_from_remaining_layers = result.add_all_annotations_from_other(
doc,
override_annotations=override_annotations,
removed_annotations=removed_annotations,
strict=strict_span_conversion,
verbose=verbose,
)
if added_annotations is not None:
for layer_name, orig_ann_id2new_ann in added_annotations_from_remaining_layers.items():
ann_id2ann = {
ann._id: ann for ann in list(doc[layer_name]) + list(doc[layer_name].predictions)
}
annotation_mapping = {
ann_id2ann[orig_ann_id]: new_ann
for orig_ann_id, new_ann in orig_ann_id2new_ann.items()
}
added_annotations.setdefault(layer_name, {}).update(annotation_mapping)
return result
def tokenize_document(
doc: TextBasedDocument,
tokenizer: PreTrainedTokenizer,
result_document_type: Type[ToD],
partition_layer: Optional[str] = None,
strict_span_conversion: bool = True,
added_annotations: Optional[List[Dict[str, Dict[Annotation, Annotation]]]] = None,
verbose: bool = True,
**tokenize_kwargs,
) -> List[ToD]:
"""Tokenize a document with a given tokenizer and return a list of token based documents. The
document is tokenized in partitions if a partition layer is provided. The annotations that
target the text are converted to target the tokens and also all dependent annotations are
converted.
Args:
doc (TextBasedDocument): The document to tokenize.
tokenizer (PreTrainedTokenizer): The tokenizer.
result_document_type (Type[ToD]): The exact type of the token based documents.
partition_layer (Optional[str], optional): The layer to use for partitioning the document. If None, the whole
document is tokenized. Defaults to None.
strict_span_conversion (bool, optional): If True, raise an error if not all annotations can be converted to
token based documents. Defaults to True.
added_annotations (Optional[List[Dict[str, Dict[Annotation, Annotation]]]], optional): Pass an empty list to
collect the added annotations. Defaults to None.
verbose (bool, optional): If True, log warnings if annotations can not be converted. Defaults to True.
Returns:
List[ToD]: The token based documents of type result_document_type with the converted annotations.
"""
added_annotation_lists: Dict[str, List[Annotation]] = defaultdict(list)
result = []
partitions: Iterable[Span]
if partition_layer is None:
partitions = [Span(start=0, end=len(doc.text))]
else:
partitions = doc[partition_layer]
for partition in partitions:
text = doc.text[partition.start : partition.end]
current_tokenize_kwargs = copy(tokenize_kwargs)
if "text" in tokenize_kwargs:
current_tokenize_kwargs["text_pair"] = text
sequence_index = 1
else:
current_tokenize_kwargs["text"] = text
sequence_index = 0
tokenized_text = tokenizer(**current_tokenize_kwargs)
for batch_encoding in tokenized_text.encodings:
token_offset_mapping = batch_encoding.offsets
char_to_token: Optional[Callable[[int], Optional[int]]]
char_to_token = functools.partial(
batch_encoding.char_to_token, sequence_index=sequence_index
)
token_offset_mapping = [
offsets if s_id == sequence_index else (0, 0)
for s_id, offsets in zip(batch_encoding.sequence_ids, token_offset_mapping)
]
if partition.start > 0:
token_offset_mapping = [
(start + partition.start, end + partition.start)
for start, end in token_offset_mapping
]
char_to_token = None
current_added_annotations: Dict[str, Dict[Annotation, Annotation]] = defaultdict(dict)
tokenized_document = text_based_document_to_token_based(
doc,
tokens=batch_encoding.tokens,
result_document_type=result_document_type,
token_offset_mapping=token_offset_mapping,
char_to_token=char_to_token,
strict_span_conversion=False,
verbose=False,
added_annotations=current_added_annotations,
)
tokenized_document.metadata["tokenizer_encoding"] = batch_encoding
result.append(tokenized_document)
for k, v in current_added_annotations.items():
added_annotation_lists[k].extend(v)
if added_annotations is not None:
added_annotations.append(current_added_annotations)
missed_annotations = defaultdict(set)
if strict_span_conversion or verbose:
# We check the annotations with respect to the layers of the result_document_type.
# Note that the original document may have more layers, but since result documents
# are of type result_document_type, we only check the layers of this type.
for annotation_field in result_document_type.annotation_fields():
# do not check the partition layer because the partitions are not required later on
# and entries get quite probably removed when windowing is applied, so this just pollutes the logs
if annotation_field.name != partition_layer:
current_missed_annotations = set(doc[annotation_field.name]) - set(
added_annotation_lists[annotation_field.name]
)
if len(current_missed_annotations) > 0:
missed_annotations[annotation_field.name] = current_missed_annotations
if len(missed_annotations) > 0:
missed_annotations_simplified = {k: str(v) for k, v in missed_annotations.items()}
if strict_span_conversion:
raise ValueError(
f"could not convert all annotations from document with id={doc.id} to token based documents, "
f"but strict_span_conversion is True, so raise an error, "
f"missed annotations:\n{json.dumps(missed_annotations_simplified, sort_keys=True, indent=2)}"
)
else:
if verbose:
logger.warning(
f"could not convert all annotations from document with id={doc.id} to token based documents, "
f"missed annotations (disable this message with verbose=False):\n"
f"{json.dumps(missed_annotations_simplified, sort_keys=True, indent=2)}"
)
return result