ModernLayoutLM / processing_modernlayoutlm.py

Upload folder using huggingface_hub

3fac549 verified 29 days ago

10.1 kB

"""
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"]

	"""
	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"]