[feat] Add UNITER text processor (#1132)

Summary:
Pull Request resolved: #1132

Add uniter_text_tokenizer which adds 'input_ids_masked'
to the sample list seperate from 'input_ids'.

Test Plan:
**Unit tests**

Tests for the construction of the processor and processor output.
We assert that the tokens are correct and contain the enhanced fields required for UNITER.

Reviewed By: ebsmothers

Differential Revision: D31865997

Pulled By: Ryan-Qiyu-Jiang

fbshipit-source-id: cf8f8b312aebe39bdfe3d5f5516831e99e27888b

Author: Ryan-Qiyu-Jiang

mmf/datasets/processors/bert_processors.py

@@ -355,6 +355,35 @@ def __init__(self, config, *args, **kwargs):
         self._probability = config.get("mask_probability", 0)
 
 
+def get_pair_text_tokens(item, masked_token_processor):
+    if "text" in item:
+        text_a = item["text"]
+    elif "text_a" in item:
+        text_a = item["text_a"]
+    else:
+        text_a = " ".join(item["tokens"])
+
+    if isinstance(text_a, list):
+        text_a = " ".join(text_a)
+
+    tokens_a = masked_token_processor.tokenize(text_a)
+
+    # 'text_b' can be defined in the dataset preparation
+    tokens_b = None
+    if "text_b" in item:
+        text_b = item["text_b"]
+        if text_b:
+            tokens_b = masked_token_processor.tokenize(text_b)
+
+    masked_token_processor._truncate_seq_pair(
+        tokens_a, tokens_b, masked_token_processor._max_seq_length
+    )
+    output = masked_token_processor._convert_to_indices(
+        tokens_a, tokens_b, probability=masked_token_processor._probability
+    )
+    return output
+
+
 @registry.register_processor("vilt_text_tokenizer")
 class VILTTextTokenizer(MaskedTokenProcessor):
     def __init__(self, config, *args, **kwargs):
@@ -372,28 +401,93 @@ def __init__(self, config, *args, **kwargs):
         self._probability = config.get("mask_probability", 0)
 
     def __call__(self, item):
-        if "text" in item:
-            text_a = item["text"]
-        elif "text_a" in item:
-            text_a = item["text_a"]
-        else:
-            text_a = " ".join(item["tokens"])
+        output = get_pair_text_tokens(item, self)
+        output["text"] = output["tokens"]
+        return output
 
-        if isinstance(text_a, list):
-            text_a = " ".join(text_a)
 
-        tokens_a = self.tokenize(text_a)
+@registry.register_processor("uniter_text_tokenizer")
+class UNITERTextTokenizer(MaskedTokenProcessor):
+    def __init__(self, config, *args, **kwargs):
+        from transformers import BertTokenizer
 
-        # 'text_b' can be defined in the dataset preparation
-        tokens_b = None
-        if "text_b" in item:
-            text_b = item["text_b"]
-            if text_b:
-                tokens_b = self.tokenize(text_b)
+        if isinstance(config, str):
+            config = {"from_pretrained": config}
 
-        self._truncate_seq_pair(tokens_a, tokens_b, self._max_seq_length)
-        output = self._convert_to_indices(
-            tokens_a, tokens_b, probability=self._probability
+        from_pretrained_name = config.get("from_pretrained", "bert-base-uncased")
+        kwargs_dict = dict(kwargs, do_lower_case="uncased" in from_pretrained_name)
+        self._tokenizer = BertTokenizer.from_pretrained(
+            from_pretrained_name, **kwargs_dict
         )
-        output["text"] = output["tokens"]
+        self._max_seq_length = config.get("max_seq_length", 25)
+        self._probability = config.get("mask_probability", 0)
+
+    def __call__(self, item: Dict[str, Any]):
+        output = get_pair_text_tokens(item, self)
+        output["text"] = output["tokens_masked"]
+        output["tokens"] = output["tokens_masked"]
+        if "is_correct" in item:
+            output["is_correct"] = torch.tensor(
+                item.get("is_correct", True), dtype=torch.long
+            )
         return output
+
+    def _token_transform(
+        self, tokens: List[str], tokens_b: Optional[List[str]] = None
+    ) -> Tuple[torch.Tensor, int, int, List[str]]:
+        tokens = [self._CLS_TOKEN] + tokens + [self._SEP_TOKEN]
+        if tokens_b:
+            tokens += tokens_b + [self._SEP_TOKEN]
+
+        input_ids = self._convert_tokens_to_ids(tokens)
+        token_len = len(input_ids)
+        token_pad = self._max_seq_length - token_len
+        # Zero-pad up to the sequence length.
+        input_ids += [self._PAD_TOKEN_ID] * token_pad
+        input_ids_tensor = torch.tensor(input_ids, dtype=torch.long)
+        return input_ids_tensor, token_len, token_pad, tokens
+
+    def _convert_to_indices(
+        self,
+        tokens_a: List[str],
+        tokens_b: Optional[List[str]] = None,
+        probability: float = 0.15,
+    ) -> Dict[str, torch.Tensor]:
+        """
+        BERT encodes
+        - single sequence: ``[CLS] X [SEP]``
+        - pair of sequences: ``[CLS] A [SEP] B [SEP]``
+        """
+        input_ids_original, _, _, _ = self._token_transform(tokens_a, tokens_b)
+
+        tokens_a, label_a = self._random_word(tokens_a, probability=probability)
+        segment_ids = [0] * (len(tokens_a) + 2)
+
+        if tokens_b:
+            tokens_b, label_b = self._random_word(tokens_b, probability=probability)
+            lm_label_ids = [-1] + label_a + [-1] + label_b + [-1]
+            assert len(tokens_b) > 0
+            segment_ids += [1] * (len(tokens_b) + 1)
+        else:
+            lm_label_ids = [-1] + label_a + [-1]
+
+        input_ids_masked, token_len, token_pad, tokens_masked = self._token_transform(
+            tokens_a, tokens_b
+        )
+
+        input_mask = [1] * token_len + [0] * token_pad
+        segment_ids += [0] * token_pad
+        lm_label_ids += [-1] * token_pad
+
+        input_mask = torch.tensor(input_mask, dtype=torch.long)
+        segment_ids = torch.tensor(segment_ids, dtype=torch.long)
+        lm_label_ids = torch.tensor(lm_label_ids, dtype=torch.long)
+        return {
+            "input_ids_masked": input_ids_masked,  # specifically for MLM heads
+            "input_ids": input_ids_original,  # unmasked tokens for CLIP heads
+            # input_mask is non-padding (1) vs padding (0) mask (not MLM token masking)
+            "input_mask": input_mask,
+            "segment_ids": segment_ids,
+            "lm_label_ids": lm_label_ids,
+            "tokens_masked": tokens_masked,
+        }

tests/datasets/test_bert_processors.py

@@ -185,3 +185,107 @@ def test_vilt_tokenizer(self):
 
         # Test [MASK] token is present
         self.assertTrue(103 in results["input_ids"])
+
+    def test_uniter_tokenizer(self):
+        from mmf.datasets.processors.bert_processors import UNITERTextTokenizer
+
+        test_utils.setup_proxy()
+        config = OmegaConf.create(
+            {
+                "tokenizer_config": {
+                    "type": "bert-base-uncased",
+                    "params": {"do_lower_case": True},
+                },
+                "mask_probability": 0.5,
+                "max_seq_length": 128,
+            }
+        )
+
+        processor = UNITERTextTokenizer(config)
+
+        # Test normal caption
+        arg = {"text": "This will be a test of tokens?"}
+        results = processor(arg)
+        expected_input_ids = torch.zeros(128, dtype=torch.long)
+        expected_input_ids[:11] = torch.tensor(
+            [101, 2023, 2097, 2022, 1037, 3231, 1997, 19204, 2015, 1029, 102],
+            dtype=torch.long,
+        )
+        expected_segment_ids = torch.zeros(128, dtype=torch.long)
+        expected_masks = torch.zeros(128, dtype=torch.long)
+        expected_masks[:11] = 1
+        self.assertTrue(torch.equal(results["input_ids"], expected_input_ids))
+        self.assertTrue(torch.equal(results["segment_ids"], expected_segment_ids))
+        self.assertTrue(torch.equal(results["input_mask"], expected_masks))
+        self.assertTrue("input_ids_masked" in results)
+        self.assertEqual(results["input_ids"].shape, results["input_ids_masked"].shape)
+
+        # Test empty caption
+        arg = {"text": ""}
+        results = processor(arg)
+        expected_input_ids = torch.zeros(128, dtype=torch.long)
+        expected_input_ids[:2] = torch.tensor([101, 102], dtype=torch.long)
+        expected_segment_ids = torch.zeros(128, dtype=torch.long)
+        expected_masks = torch.zeros(128, dtype=torch.long)
+        expected_masks[:2] = 1
+        self.assertTrue(torch.equal(results["input_ids"], expected_input_ids))
+        self.assertTrue(torch.equal(results["segment_ids"], expected_segment_ids))
+        self.assertTrue(torch.equal(results["input_mask"], expected_masks))
+        self.assertTrue("input_ids_masked" in results)
+        self.assertEqual(results["input_ids"].shape, results["input_ids_masked"].shape)
+
+        # Test long caption
+        arg = {"text": "I am working for facebook " * 100}  # make a long sentence
+        results = processor(arg)
+        expected_input_ids = [1045, 2572, 2551, 2005, 9130] * 100
+        expected_input_ids.insert(0, 101)  # [CLS]
+        expected_input_ids = expected_input_ids[:128]
+        expected_input_ids[-1] = 102  # [SEP]
+        expected_input_ids = torch.tensor(expected_input_ids, dtype=torch.long)
+        expected_segment_ids = torch.zeros(128, dtype=torch.long)
+        expected_masks = torch.ones(128, dtype=torch.long)
+        self.assertTrue(torch.equal(results["input_ids"], expected_input_ids))
+        self.assertTrue(torch.equal(results["segment_ids"], expected_segment_ids))
+        self.assertTrue(torch.equal(results["input_mask"], expected_masks))
+        self.assertTrue("input_ids_masked" in results)
+        self.assertEqual(results["input_ids"].shape, results["input_ids_masked"].shape)
+
+        # Test two captions
+        arg = {
+            "text_a": "This will be a test of tokens?",
+            "text_b": "I am working for facebook",
+        }
+        results = processor(arg)
+        expected_input_ids = torch.zeros(128, dtype=torch.long)
+        expected_input_ids[:17] = torch.tensor(
+            [101, 2023, 2097, 2022, 1037, 3231, 1997, 19204, 2015, 1029, 102]
+            + [1045, 2572, 2551, 2005, 9130, 102],
+            dtype=torch.long,
+        )
+        expected_segment_ids = torch.zeros(128, dtype=torch.long)
+        expected_segment_ids[11:17] = 1
+        expected_masks = torch.zeros(128, dtype=torch.long)
+        expected_masks[:17] = 1
+        self.assertTrue(torch.equal(results["input_ids"], expected_input_ids))
+        self.assertTrue(torch.equal(results["segment_ids"], expected_segment_ids))
+        self.assertTrue(torch.equal(results["input_mask"], expected_masks))
+        self.assertTrue("input_ids_masked" in results)
+        self.assertEqual(results["input_ids"].shape, results["input_ids_masked"].shape)
+
+        # Test masked caption
+        processor._probability = 1.0
+        arg = {"text": "This will be a test of tokens?"}
+        results = processor(arg)
+        expected_input_ids = torch.zeros(128, dtype=torch.long)
+        expected_input_ids[:11] = torch.tensor(
+            [101, 2023, 2097, 2022, 1037, 3231, 1997, 19204, 2015, 1029, 102],
+            dtype=torch.long,
+        )
+        expected_segment_ids = torch.zeros(128, dtype=torch.long)
+        self.assertTrue(torch.equal(results["input_ids"], expected_input_ids))
+        self.assertTrue(torch.equal(results["segment_ids"], expected_segment_ids))
+        self.assertTrue("input_ids_masked" in results)
+        self.assertEqual(results["input_ids"].shape, results["input_ids_masked"].shape)
+
+        # Test [MASK] token is present
+        self.assertTrue(103 in results["input_ids_masked"])