open source transformer based models - data, tensorizers and tokenizer (facebookresearch#708)

Summary:
Pull Request resolved: facebookresearch#708

Open source BERTTensorizer, XLMTensorizer, WordpieceTokenizer and PackedLMData.  Working up to open sourcing MaskedLM and BERT classification models.

Reviewed By: rutyrinott

Differential Revision: D15868227

fbshipit-source-id: 891730437b570d5b152ab2da4ea19c4aaef6bb2b

Author: borguz

pytext/data/bert_tensorizer.py

@@ -0,0 +1,91 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+import itertools
+from typing import List
+
+from fairseq.data.masked_lm_dictionary import BertDictionary
+from pytext.config.component import ComponentType, create_component
+from pytext.data.tensorizers import TokenTensorizer
+from pytext.data.tokenizers import Tokenizer, WordPieceTokenizer
+from pytext.data.utils import BOS, EOS, MASK, PAD, UNK, Vocabulary, pad_and_tensorize
+
+
+class BERTTensorizer(TokenTensorizer):
+    """
+    Tensorizer for BERT tasks.  Works for single sentence, sentence pair, triples etc.
+    """
+
+    __EXPANSIBLE__ = True
+
+    class Config(TokenTensorizer.Config):
+        #: The tokenizer to use to split input text into tokens.
+        columns: List[str] = ["text"]
+        tokenizer: Tokenizer.Config = WordPieceTokenizer.Config()
+        add_bos_token: bool = False
+        add_eos_token: bool = True
+        bos_token: str = "[CLS]"
+        eos_token: str = "[SEP]"
+        pad_token: str = "[PAD]"
+        unk_token: str = "[UNK]"
+        mask_token: str = "[MASK]"
+        vocab_file: str = WordPieceTokenizer.Config().wordpiece_vocab_path
+
+    @classmethod
+    def from_config(cls, config: Config, **kwargs):
+        tokenizer = create_component(ComponentType.TOKENIZER, config.tokenizer)
+        replacements = {
+            config.unk_token: UNK,
+            config.pad_token: PAD,
+            config.bos_token: BOS,
+            config.eos_token: EOS,
+            config.mask_token: MASK,
+        }
+        if isinstance(tokenizer, WordPieceTokenizer):
+            vocab = Vocabulary(
+                [token for token, _ in tokenizer.vocab.items()],
+                replacements=replacements,
+            )
+        else:
+            dictionary = BertDictionary.load(config.vocab_file)
+            vocab = Vocabulary(
+                dictionary.symbols, dictionary.count, replacements=replacements
+            )
+        return cls(
+            columns=config.columns,
+            tokenizer=tokenizer,
+            add_bos_token=config.add_bos_token,
+            add_eos_token=config.add_eos_token,
+            use_eos_token_for_bos=config.use_eos_token_for_bos,
+            max_seq_len=config.max_seq_len,
+            vocab=vocab,
+            **kwargs,
+        )
+
+    def __init__(self, columns, **kwargs):
+        super().__init__(text_column=None, **kwargs)
+        self.columns = columns
+        # Manually initialize column_schema since we are sending None to TokenTensorizer
+        self.column_schema = [(column, str) for column in columns]
+
+    def numberize(self, row):
+        """Tokenize, look up in vocabulary."""
+        sentences = [self._lookup_tokens(row[column])[0] for column in self.columns]
+        sentences[0] = [self.vocab.idx[BOS]] + sentences[0]
+        seq_lens = (len(sentence) for sentence in sentences)
+        segment_labels = ([i] * seq_len for i, seq_len in enumerate(seq_lens))
+        tokens = list(itertools.chain(*sentences))
+        segment_labels = list(itertools.chain(*segment_labels))
+        seq_len = len(tokens)
+        # tokens, segment_label, seq_len
+        return tokens, segment_labels, seq_len
+
+    def sort_key(self, row):
+        return row[2]
+
+    def tensorize(self, batch):
+        tokens, segment_labels, seq_lens = zip(*batch)
+        tokens = pad_and_tensorize(tokens, self.vocab.get_pad_index())
+        pad_mask = (tokens != self.vocab.get_pad_index()).long()
+        segment_labels = pad_and_tensorize(segment_labels, self.vocab.get_pad_index())
+        return tokens, pad_mask, segment_labels

pytext/data/packed_lm_data.py

@@ -0,0 +1,149 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+from typing import Dict, List, Optional, Type
+
+from pytext.common.constants import Stage
+from pytext.data import Batcher, Data
+from pytext.data.bert_tensorizer import BERTTensorizer
+from pytext.data.data import RowData
+from pytext.data.sources import DataSource
+from pytext.data.tensorizers import Tensorizer, TokenTensorizer
+from pytext.data.xlm_tensorizer import XLMTensorizer
+
+
+class PackedLMData(Data):
+    """
+    Special purpose Data object which assumes a single text tensorizer.  Packs
+    tokens into a square batch with no padding.  Used for LM training. The object
+    also takes in an optional language argument which is used for cross-lingual
+    LM training.
+    """
+
+    __EXPANSIBLE__ = True
+
+    class Config(Data.Config):
+        max_seq_len: int = 128
+
+    @classmethod
+    def from_config(
+        cls,
+        config: Config,
+        schema: Dict[str, Type],
+        tensorizers: Dict[str, Tensorizer],
+        language: Optional[str] = None,
+        rank: int = 0,
+        world_size: int = 1,
+    ):
+        return super(PackedLMData, cls).from_config(
+            config,
+            schema,
+            tensorizers,
+            rank,
+            world_size,
+            language=language,
+            max_seq_len=config.max_seq_len,
+        )
+
+    def __init__(
+        self,
+        data_source: DataSource,
+        tensorizers: Dict[str, Tensorizer],
+        batcher: Batcher = None,
+        max_seq_len: int = Config.max_seq_len,
+        sort_key: Optional[str] = None,
+        # language is used in cross-lingual LM training
+        language: Optional[str] = None,
+        in_memory: Optional[bool] = False,
+    ):
+        super().__init__(data_source, tensorizers, batcher, sort_key, in_memory)
+        assert len(list(self.tensorizers.items())) == 1
+        self.tensorizer_name, self.tensorizer = list(self.tensorizers.items())[0]
+        self.remainder: Dict[str, List[int]] = {"tokens": [], "segment_labels": []}
+        self.max_seq_len = max_seq_len
+        self.language = language
+        self.batch = {Stage.TRAIN: None, Stage.EVAL: None, Stage.TEST: None}
+
+    def _parse_row(self, row):
+        """
+        The output of numberization has different number of elements depending on
+        the tensorizer used. For example: positions tensor is only output by the
+        XLMTensorizer. This function unpacks the elements according to the
+        specific tensorizer used.
+        Additionally, since we are packing tokens into fixed size
+        blocks, we don't need to use the positions vector output by the call to
+        numberize. We will simply create this in `_format_output_row`.
+        """
+        numberized_row = self.tensorizer.numberize(row)
+        if isinstance(self.tensorizer, XLMTensorizer):
+            tokens, seq_len, segment_labels, _ = numberized_row
+        elif isinstance(self.tensorizer, BERTTensorizer):
+            tokens, segment_labels, seq_len = numberized_row
+        elif isinstance(self.tensorizer, TokenTensorizer):
+            tokens, seq_len, _ = numberized_row
+            segment_labels = []
+        else:
+            raise NotImplementedError(
+                "PackedLMData only supports XLMTensorizer, BERTTensorizer and "
+                "TokenTensorizer."
+            )
+        return tokens, segment_labels, seq_len
+
+    def _format_output_row(self, tokens, segment_labels, seq_len):
+        """
+        The tensorize function for different tensorizers takes in different
+        number of inputs which may be arranged differently. This function formats
+        the output dict to conform to the expectations of the tensorizer.
+        In case of the XLMTensorizer, we also need to create a new positions list
+        which goes from 0 to seq_len.
+        """
+        if isinstance(self.tensorizer, XLMTensorizer):
+            positions = [index for index in range(seq_len)]
+            return {self.tensorizer_name: (tokens, seq_len, segment_labels, positions)}
+        elif isinstance(self.tensorizer, BERTTensorizer):
+            return {self.tensorizer_name: (tokens, segment_labels, seq_len)}
+        elif isinstance(self.tensorizer, TokenTensorizer):
+            # dummy token_ranges
+            return {self.tensorizer_name: (tokens, seq_len, [(-1, -1)] * seq_len)}
+        else:
+            raise NotImplementedError(
+                "PackedLMData only supports BERTTensorizer and TokenTensorizer."
+            )
+
+    def _yield_and_reset(self):
+        packed_tokens = list(self.remainder["tokens"])
+        packed_segments = list(self.remainder["segment_labels"])
+        self.remainder: Dict[str, List[int]] = {"tokens": [], "segment_labels": []}
+        return RowData(
+            {},  # packed LM data doesn't respect data cardinality
+            self._format_output_row(packed_tokens, packed_segments, len(packed_tokens)),
+        )
+
+    def numberize_rows(self, rows):
+        """
+        This function does the actual packing. It processes rows until we obtain
+        a block of data with length = max_seq_len.
+        """
+        for row in rows:
+
+            # if the packedLM object has a language member then a cross-lingual
+            # LM is being trained using monolingual data.
+            # Add this language to the row since the underlying
+            # tensorizer needs this to generate language embeddings (used as
+            # segment_labels below)
+            if self.language:
+                row["language"] = self.language
+
+            tokens, segment_labels, seq_len = self._parse_row(row)
+            remaining = self.max_seq_len - len(self.remainder["tokens"]) - 1
+            while remaining < len(tokens):
+                self.remainder["tokens"].extend(tokens[:remaining])
+                self.remainder["segment_labels"].extend(segment_labels[:remaining])
+                tokens = tokens[remaining:]
+                segment_labels = segment_labels[remaining:]
+                yield self._yield_and_reset()
+                remaining = self.max_seq_len - 1
+            self.remainder["tokens"].extend(tokens)
+            self.remainder["segment_labels"].extend(segment_labels)
+        if len(self.remainder["tokens"]):
+            yield self._yield_and_reset()