File size: 10,109 Bytes
3fac549 | 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 | """
Turning words and boxes into model inputs
There is no image tower, so this only tokenizes words and broadcasts each box across its subwords
This module is copied into exported model directories and may import only torch and transformers
"""
from __future__ import annotations
from typing import Any
from transformers.processing_utils import ProcessorMixin
from transformers.tokenization_utils_base import BatchEncoding
# Follow LayoutLM: CLS at the origin, SEP at the far corner, padding masked out
CLS_BOX = (0, 0, 0, 0)
SEP_BOX = (1000, 1000, 1000, 1000)
PAD_BOX = (0, 0, 0, 0)
# Label id meaning "do not compute loss here", the torch cross-entropy default
IGNORE_INDEX = -100
class ModernLayoutLmProcessor(ProcessorMixin):
"""
Pairs a fast tokenizer with box and label expansion over subwords
Args:
tokenizer: Any fast tokenizer, so the LayoutLM baseline reuses this with its own one.
coordinate_scale: The scale boxes are expected to already be normalized to.
label_all_subwords: Label every subword, not only the first as published FUNSD and CORD do.
"""
attributes = ["tokenizer"]
tokenizer_class = "AutoTokenizer"
def __init__(
self,
tokenizer: Any,
coordinate_scale: int = 1000,
label_all_subwords: bool = False,
) -> None:
self.coordinate_scale = coordinate_scale
self.label_all_subwords = label_all_subwords
super().__init__(tokenizer)
# Diverge as ProcessorMixin.__call__ is typed for image and audio processors
def __call__( # pyright: ignore[reportIncompatibleMethodOverride]
self,
words: list[str] | list[list[str]],
boxes: list[list[int]] | list[list[list[int]]],
word_labels: list[int] | list[list[int]] | None = None,
max_length: int | None = None,
padding: bool | str = "max_length",
truncation: bool = True,
return_overflowing_tokens: bool = False,
stride: int = 0,
return_tensors: str | None = None,
**kwargs: Any,
) -> BatchEncoding:
"""
Tokenizes words and aligns boxes, and optionally labels, to the resulting subwords
Accepts one document or a batch, told apart by whether `words` holds strings or lists.
Special tokens receive the sentinel boxes defined in this module.
With `return_overflowing_tokens` a long page becomes several chunks, each with its boxes.
Args:
words: One document's words, or a batch of documents' words.
boxes: One (x0, y0, x1, y1) box per word on `coordinate_scale`, batched like `words`.
word_labels: One class index per word, batched like `words`, or None for no labels.
max_length: Truncation length, by default the tokenizer's own.
padding: Passed to the tokenizer.
truncation: Passed to the tokenizer.
return_overflowing_tokens: Emit every chunk of an over-length page, not only the first.
stride: Overlap in tokens between consecutive chunks.
return_tensors: "pt" for torch tensors, None for lists.
**kwargs: Forwarded to the tokenizer.
Returns:
A BatchEncoding of `input_ids`, `attention_mask`, `bbox`, plus `labels` if given any.
Raises:
ValueError: If `words` and `boxes` differ in length, or a coordinate leaves the grid.
"""
batched = bool(words) and isinstance(words[0], list)
batch_words: list[list[str]] = words if batched else [words] # type: ignore[assignment]
batch_boxes: list[list[list[int]]] = boxes if batched else [boxes] # type: ignore[assignment]
batch_labels: list[list[int]] | None = None
if word_labels is not None:
batch_labels = word_labels if batched else [word_labels] # type: ignore[assignment]
self._validate(batch_words, batch_boxes, batch_labels)
encoding = self.tokenizer(
text=batch_words,
is_split_into_words=True,
max_length=max_length,
padding=padding,
truncation=truncation,
return_overflowing_tokens=return_overflowing_tokens,
stride=stride,
return_tensors=None,
**kwargs,
)
sample_map = encoding.get("overflow_to_sample_mapping")
if sample_map is None:
sample_map = list(range(len(encoding["input_ids"])))
aligned_boxes: list[list[list[int]]] = []
aligned_labels: list[list[int]] = []
for chunk_index, sample_index in enumerate(sample_map):
word_ids = encoding.word_ids(chunk_index)
source_boxes = batch_boxes[sample_index]
aligned_boxes.append(self._expand_boxes(word_ids, source_boxes))
if batch_labels is not None:
aligned_labels.append(self._expand_labels(word_ids, batch_labels[sample_index]))
encoding["bbox"] = aligned_boxes
if batch_labels is not None:
encoding["labels"] = aligned_labels
if return_tensors is not None:
encoding.convert_to_tensors(tensor_type=return_tensors)
return encoding
def _validate(
self,
batch_words: list[list[str]],
batch_boxes: list[list[list[int]]],
batch_labels: list[list[int]] | None,
) -> None:
"""
Checks that the parallel inputs line up and the coordinates are in range
A misalignment would otherwise train a model on wrong boxes without ever erroring.
Args:
batch_words: Words per document.
batch_boxes: Boxes per document.
batch_labels: Labels per document, or None.
Raises:
ValueError: On a length mismatch or an out of range coordinate.
"""
if len(batch_words) != len(batch_boxes):
raise ValueError(f"Got {len(batch_words)} word lists but {len(batch_boxes)} box lists")
if batch_labels is not None and len(batch_labels) != len(batch_words):
raise ValueError(
f"Got {len(batch_words)} word lists but {len(batch_labels)} label lists"
)
for index, (doc_words, doc_boxes) in enumerate(zip(batch_words, batch_boxes, strict=True)):
if len(doc_words) != len(doc_boxes):
raise ValueError(
f"Document {index} has {len(doc_words)} words but {len(doc_boxes)} boxes"
)
if batch_labels is not None and len(batch_labels[index]) != len(doc_words):
raise ValueError(
f"Document {index} has {len(doc_words)} words "
f"but {len(batch_labels[index])} labels"
)
for box in doc_boxes:
if len(box) != 4:
raise ValueError(f"Document {index} has a box with {len(box)} coordinates")
if any(not 0 <= value <= self.coordinate_scale for value in box):
raise ValueError(
f"Document {index} has box {box} outside "
f"[0, {self.coordinate_scale}]. Normalize before calling the processor."
)
def _expand_boxes(
self,
word_ids: list[int | None],
source_boxes: list[list[int]],
) -> list[list[int]]:
"""
Assigns a box to every token position in one chunk
Args:
word_ids: Per-token index into `source_boxes`, None for special and padding tokens.
source_boxes: One box per word of the source document.
Returns:
One box per token position.
"""
special_boxes = self._special_boxes(word_ids)
return [
list(source_boxes[word_id]) if word_id is not None else special_boxes[position]
for position, word_id in enumerate(word_ids)
]
def _special_boxes(self, word_ids: list[int | None]) -> dict[int, list[int]]:
"""
Chooses a sentinel box for each token that came from no word
Leading specials sit at the origin like CLS, and the first trailing one at the far corner.
Everything past it is padding at the origin.
Args:
word_ids: Per-token index into the source words, None for non-word tokens.
Returns:
A mapping from token position to sentinel box.
"""
word_positions = [i for i, word_id in enumerate(word_ids) if word_id is not None]
last_word = word_positions[-1] if word_positions else -1
boxes: dict[int, list[int]] = {}
seen_separator = False
for position, word_id in enumerate(word_ids):
if word_id is not None:
continue
if position < last_word:
boxes[position] = list(CLS_BOX)
elif not seen_separator:
boxes[position] = list(SEP_BOX)
seen_separator = True
else:
boxes[position] = list(PAD_BOX)
return boxes
def _expand_labels(
self,
word_ids: list[int | None],
source_labels: list[int],
) -> list[int]:
"""
Assigns a label to every token position in one chunk
Special and padding tokens always get `IGNORE_INDEX`.
Whether the later subwords of a word inherit its label follows `label_all_subwords`.
Args:
word_ids: Per-token index into `source_labels`, None for non-word tokens.
source_labels: One class index per word of the source document.
Returns:
One label per token position.
"""
labels: list[int] = []
previous = None
for word_id in word_ids:
if word_id is None:
labels.append(IGNORE_INDEX)
elif word_id != previous or self.label_all_subwords:
labels.append(source_labels[word_id])
else:
labels.append(IGNORE_INDEX)
previous = word_id
return labels
__all__ = ["ModernLayoutLmProcessor"]
|