src/consec_tokenizer.py

from functools import lru_cache
from typing import List, NamedTuple, Optional, Dict, Tuple, Union

import torch
from tokenizers import AddedToken

from transformers import AutoTokenizer, PreTrainedTokenizer, PreTrainedTokenizerFast

from src.utils.collections import flatten


class TokenizationOutput(NamedTuple):
    input_ids: torch.Tensor
    attention_mask: torch.Tensor
    token_type_ids: Optional[torch.Tensor]
    definitions_offsets: Dict[str, Tuple[int, int]]
    relative_positions: Optional[torch.Tensor] = None


# utilities to extract the gold definition indices in the input sequence
def extract_gold_indices(
    definition: str, definitions_sequence: str, offsets_mapping: List[Tuple[int, int]]
) -> Tuple[int, int]:
    gold_start_offset = definitions_sequence.index(definition)
    gold_end_offset = gold_start_offset + len(definition)

    offset_start2bpe_idx = {}
    offset_end2bpe_idx = {}
    for i, (off_start, off_end) in enumerate(offsets_mapping):

        if off_start == off_end:
            continue  # specials

        offset_start2bpe_idx[off_start] = i
        offset_end2bpe_idx[off_end] = i

    start_bpe_idx = offset_start2bpe_idx[gold_start_offset]
    end_bpe_idx = offset_end2bpe_idx[gold_end_offset]

    assert start_bpe_idx < end_bpe_idx

    return start_bpe_idx, end_bpe_idx


class ConsecTokenizer:
    def __init__(
        self,
        transformer_model: Union[str, PreTrainedTokenizer, PreTrainedTokenizerFast],
        target_marker: Tuple[str, str],
        context_definitions_token: str,
        context_markers: Dict,
        add_prefix_space: bool,
    ):
        if type(transformer_model) == str:
            self.tokenizer = AutoTokenizer.from_pretrained(transformer_model)
        else:
            self.tokenizer = transformer_model
        self.target_marker = target_marker
        self.context_definitions_token = context_definitions_token
        self.context_markers = [
            (context_markers["pattern"][0].replace("#I#", f"{i}"), context_markers["pattern"][1].replace("#I#", f"{i}"))
            for i in range(context_markers["number"])
        ]
        assert (
            len(set(self.context_markers)) == context_markers["number"]
        ), f"Error in given pattern: number of unique created patterns != specified number"

        additional_special_tokens = [
            *[AddedToken(t, single_word=True, lstrip=True) for t in self.target_marker],
            *[AddedToken(t, single_word=True, lstrip=True) for p in self.context_markers for t in p],
            AddedToken(context_definitions_token, single_word=True, lstrip=True),
        ]
        self.tokenizer.add_special_tokens({"additional_special_tokens": additional_special_tokens})
        self.n_added_special_tokens = len(additional_special_tokens)
        self.add_prefix_space = add_prefix_space

    def mark_token(self, token: str, marker: Tuple[str, str]) -> str:
        bom, eom = marker
        if self.add_prefix_space:
            return f"{bom} {token} {eom}"
        else:
            return f"{bom}{token}{eom}"

    def tokenize(
        self,
        sentence: List[str],  # tokens can be multiword
        instance_idx: int,
        instance_possible_definitions: List[str],
        context_definitions2positions: List[Tuple[str, int]],
    ) -> TokenizationOutput:
        raise NotImplementedError

    @property
    def pad_token_id(self) -> int:
        return self.tokenizer.pad_token_id

    @property
    def model_max_length(self) -> int:
        return self.tokenizer.model_max_length


class DeBERTaTokenizer(ConsecTokenizer):
    def __init__(
        self,
        transformer_model: str,
        target_marker: Tuple[str, str],
        context_definitions_token: str,
        context_markers: Dict,
        add_prefix_space: bool,
        optimize_relative_positions: bool = True,
        enforce_symmetry: bool = True,
    ):
        super().__init__(transformer_model, target_marker, context_definitions_token, context_markers, add_prefix_space)
        self.optimize_relative_positions = optimize_relative_positions
        self.enforce_symmetry = enforce_symmetry

    def tokenize(
        self,
        sentence: List[str],  # tokens can be multiword
        instance_idx: int,
        instance_possible_definitions: List[str],
        context_definitions2positions: List[Tuple[str, int]],
        **kwargs,
    ) -> TokenizationOutput:
        if self.optimize_relative_positions:
            return self.power_tokenize(
                sentence, instance_idx, instance_possible_definitions, context_definitions2positions
            )
        else:
            return self.plain_tokenize(
                " ".join(sentence),
                instance_possible_definitions,
                [x[0] for x in context_definitions2positions],
                **kwargs,
            )

    def deberta_tokenize(self, text: str) -> List[int]:
        return self.tokenizer(text, return_attention_mask=False, return_token_type_ids=False, add_special_tokens=True,)[
            "input_ids"
        ][1:-1]

    def plain_tokenize(
        self,
        sentence: Union[str, List[int]],
        instance_possible_definitions: Union[List[str], List[Tuple[str, List[int]]]],
        context_definitions: Union[List[str], List[Tuple[str, List[int]]]],
        use_specials: bool = True,
    ) -> TokenizationOutput:
        if type(sentence) != list:
            sentence = f" {sentence}"
            sentence_input_ids = self.deberta_tokenize(sentence)
        else:
            sentence_input_ids = sentence

        final_input_ids = [self.tokenizer.cls_token_id] + sentence_input_ids + [self.tokenizer.sep_token_id]
        token_type_ids = [0] * len(final_input_ids)

        definitions_offsets = dict()
        for definition in instance_possible_definitions:
            if type(definition) == tuple:
                definition, definition_ids = definition
            else:
                definition_ids = self.deberta_tokenize(f" {definition}")

            definitions_offsets[definition] = len(final_input_ids), len(final_input_ids) + len(definition_ids)

            final_input_ids += definition_ids
            token_type_ids += [1] * len(definition_ids)

        # "context" definitions token
        if self.context_definitions_token is not None and use_specials:
            final_input_ids += self.deberta_tokenize(self.context_definitions_token)
            token_type_ids.append(1)
        else:
            if use_specials:
                assert len(context_definitions) == 0

        for definition in context_definitions:
            if type(definition) == tuple:
                definition, definition_ids = definition
            else:
                definition_ids = self.deberta_tokenize(f" {definition}")

            final_input_ids += definition_ids
            token_type_ids += [1] * len(definition_ids)

        # last [SEP] token
        final_input_ids.append(self.tokenizer.sep_token_id)
        token_type_ids.append(1)

        final_input_ids = torch.tensor(final_input_ids, dtype=torch.long)
        attention_mask = torch.ones_like(final_input_ids)
        token_type_ids = torch.tensor(token_type_ids, dtype=torch.long)

        return TokenizationOutput(final_input_ids, attention_mask, token_type_ids, definitions_offsets)

    @lru_cache(maxsize=10_000)
    def split_tokenize(self, word: str):
        return self.deberta_tokenize(f" {word.strip()}")

    @lru_cache(maxsize=10_000)
    def _tokenize_sentence_list(self, sentence: str) -> Tuple[List[int], List[Tuple[int, int]]]:
        sentence = sentence.split("\t")
        sentence_input_ids = [self.tokenizer.cls_token_id]
        token_offsets = []

        for token in sentence:
            token_ids = self.split_tokenize(f" {token.strip()}")
            token_offsets.append((len(sentence_input_ids), len(sentence_input_ids) + len(token_ids)))
            sentence_input_ids += token_ids

        sentence_input_ids.append(self.tokenizer.sep_token_id)

        return sentence_input_ids, token_offsets

    @staticmethod
    def mirror_arange(size: int, zero_pos: int):
        return torch.cat([torch.flip(torch.arange(0, zero_pos + 1), dims=(0,)), -torch.arange(1, size - zero_pos)])

    def power_tokenize(
        self,
        sentence: List[str],  # tokens can be multiword
        instance_idx: int,
        instance_possible_definitions: List[str],
        context_definitions2positions: List[Tuple[str, int]],
    ) -> TokenizationOutput:

        sentence_input_ids, token_offsets = self._tokenize_sentence_list("\t".join(sentence))

        instance_possible_definitions_ids = [self.deberta_tokenize(f" {ipd}") for ipd in instance_possible_definitions]
        context_definitions_ids = [self.deberta_tokenize(f" {cd}") for cd, _ in context_definitions2positions]

        total_input_ids = (
            len(sentence_input_ids)
            + len(flatten(instance_possible_definitions_ids))
            + len(flatten(context_definitions_ids))
        )

        relative_positions = torch.zeros((total_input_ids + 1, total_input_ids + 1), dtype=torch.long)

        # filling the sentences relative positions
        for id_idx in range(len(sentence_input_ids)):
            relative_positions[id_idx, :-1] = torch.cat(
                [self.mirror_arange(len(sentence_input_ids), id_idx)]
                + [
                    -torch.arange(len(sentence_input_ids) - id_idx, len(sentence_input_ids) + len(def_ids) - id_idx)
                    for def_ids in instance_possible_definitions_ids + context_definitions_ids
                ]
            )

        # flipping the matrix
        relative_positions.T[: len(sentence_input_ids)] = -relative_positions[: len(sentence_input_ids)]

        curr_offset = len(sentence_input_ids)

        # computing the relative positions for the possible definitions
        definitions2positions = [(instance_idx, ipdid) for ipdid in instance_possible_definitions_ids]
        definitions2positions += [
            (rel_idx, cdid) for cdid, (_, rel_idx) in zip(context_definitions_ids, context_definitions2positions)
        ]
        for def_num in range(len(definitions2positions)):

            rel_token_idx, def_ids = definitions2positions[def_num]
            possible_def_token_pos = list(range(len(def_ids)))
            instance_start_pos, instance_end_pos = token_offsets[rel_token_idx]

            for token_pos in possible_def_token_pos:

                for off_idx, inst_token_pos in enumerate(range(instance_start_pos, instance_end_pos)):
                    relative_positions[curr_offset + token_pos, inst_token_pos] = (
                        token_pos + (instance_end_pos - instance_start_pos) - off_idx
                    )
                    if self.enforce_symmetry or rel_token_idx != instance_idx:
                        relative_positions[inst_token_pos, curr_offset + token_pos] = -relative_positions[
                            curr_offset + token_pos, inst_token_pos
                        ]

                relative_positions[curr_offset + token_pos, len(sentence_input_ids) : -1] = self.mirror_arange(
                    total_input_ids - len(sentence_input_ids), curr_offset + token_pos - len(sentence_input_ids)
                )

            curr_offset += len(possible_def_token_pos)

        last_token_positions = torch.min(relative_positions, dim=-1)[0] - 1
        relative_positions[-1] = -last_token_positions
        relative_positions.T[-1] = last_token_positions
        relative_positions[-1, -1] = 0

        relative_positions[relative_positions == -0] = 0

        tokenization_output = self.plain_tokenize(
            sentence_input_ids[1:-1],
            [(x, y) for x, y in zip(instance_possible_definitions, instance_possible_definitions_ids)],
            [(x, y) for x, y in zip([x[0] for x in context_definitions2positions], context_definitions_ids)],
            use_specials=False,
        )

        assert tokenization_output.input_ids.shape[0] == relative_positions.shape[0]

        return TokenizationOutput(
            tokenization_output.input_ids,
            tokenization_output.attention_mask,
            tokenization_output.token_type_ids,
            tokenization_output.definitions_offsets,
            relative_positions,
        )

    @property
    def model_max_length(self) -> int:
        return 24_528  # from paper: deberta-large


class MBartTokenizer(ConsecTokenizer):
    def __init__(
        self,
        transformer_model: str,
        target_marker: Tuple[str, str],
        context_definitions_token: str,
        context_markers: Dict,
        add_prefix_space: bool,
        source_language: str = "en_XX",
        target_language: str = "en_EN",
    ):
        tokenizer = AutoTokenizer.from_pretrained(transformer_model, src_lang=source_language, tgt_lang=target_language)
        self.source_language = source_language
        super().__init__(tokenizer, target_marker, context_definitions_token, context_markers, add_prefix_space)

    def mbart_tokenize(self, text: str) -> List[int]:
        tokenization_out = self.tokenizer(
            text,
            return_attention_mask=False,
            return_token_type_ids=False,
            add_special_tokens=True,
        )["input_ids"]
        return tokenization_out[:-2]

    def tokenize(
        self,
        sentence: Union[str, List[int]],
        instance_idx: int,
        instance_possible_definitions: List[str],
        context_definitions2positions: List[Tuple[str, int]],
    ) -> TokenizationOutput:
        use_specials = False
        context_definitions = [x[0] for x in context_definitions2positions]

        sentence = f" {' '.join(sentence)}"
        sentence_input_ids = self.mbart_tokenize(sentence)

        final_input_ids = sentence_input_ids

        definitions_offsets = dict()
        for definition in instance_possible_definitions:
            if type(definition) == tuple:
                definition, definition_ids = definition
            else:
                definition_ids = self.mbart_tokenize(f" {definition}")

            definitions_offsets[definition] = len(final_input_ids), len(final_input_ids) + len(definition_ids)

            final_input_ids += definition_ids

        # "context" definitions token
        if self.context_definitions_token is not None and use_specials:
            final_input_ids += self.mbart_tokenize(self.context_definitions_token)
        else:
            if use_specials:
                assert len(context_definitions) == 0

        for definition in context_definitions:
            if type(definition) == tuple:
                definition, definition_ids = definition
            else:
                definition_ids = self.mbart_tokenize(f" {definition}")

            final_input_ids += definition_ids

        # last [SEP] token
        final_input_ids.append(self.tokenizer.sep_token_id)

        # lang id
        final_input_ids += self.mbart_tokenize(self.source_language)

        final_input_ids = torch.tensor(final_input_ids, dtype=torch.long)
        attention_mask = torch.ones_like(final_input_ids)

        return TokenizationOutput(final_input_ids, attention_mask, None, definitions_offsets)