Repository: huggingface/transformers

Issue:
FillMaskPipeline very slow when provided with a large `targets`
## Environment info
<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->

- `transformers` version: 4.6.1
- Platform: Linux-5.4.0-67-generic-x86_64-with-glibc2.10
- Python version: 3.8.5
- PyTorch version (GPU?): 1.8.1 (False)
- Tensorflow version (GPU?): N/A
- Using GPU in script?: No
- Using distributed or parallel set-up in script?: No

### Who can help
@LysandreJik @Narsil 

## Information

The model I am using: `ethanyt/guwenbert-base`, with a `RoBERTa` model and a `BertTokenizerFast` tokenizer.

## To reproduce

Steps to reproduce the behavior:

1. Initialize a `fill-mask` pipeline with the model and the tokenizer mentioned above
2. Call it with any sentence and a large `targets` (with a length of ~10k single words)

## Problem

The call would be much slower than a similar call without a `targets` argument. A call without a `targets` argument costs ~0.1s, while a call with a `targets` argument costs ~0.3s.

The following code is present in `src/transformers/pipelines/fill_mask.py`:

```python
class FillMaskPipeline(Pipeline):
    # ...
    def __call__(self, *args, targets=None, top_k: Optional[int] = None, **kwargs):
        # ...
        if targets is not None:
            # ...
            targets_proc = []
            for target in targets:
                target_enc = self.tokenizer.tokenize(target)
                # ...
                targets_proc.append(target_enc[0])
```

This function iterates through targets, rather than sending it directly to `tokenize`, which does not utilize the batch processing optimization of `TokenizerFast`s, hence the slow speed.


Discussion:
Do you have an example to reproduce the issue ? Benchmarking is sometimes suprising and hardware dependant.

I can imagine that this is indeed a slowdown as Python - Rust communication is not free.
However the omitted part of your comment is error detection, which is important and we need to keep it.

```python
            if len(targets) == 0 or len(targets[0]) == 0:
                raise ValueError("At least one target must be provided when passed.")
            if isinstance(targets, str):
                targets = [targets]

            targets_proc = []
            for target in targets:
                target_enc = self.tokenizer.tokenize(target)
                if len(target_enc) > 1 or target_enc[0] == self.tokenizer.unk_token:
                    logger.warning(
                        f"The specified target token `{target}` does not exist in the model vocabulary. "
                        f"Replacing with `{target_enc[0]}`."
                    )
                targets_proc.append(target_enc[0])
            target_inds = np.array(self.tokenizer.convert_tokens_to_ids(targets_proc))
```

I think we can get away with encoding every target at once, then iterating through the whole array to do the error detection.

However, as this is a performance problem, I think we realistically need to test that improving performance on 10K targets, does not reduce performance significantly on 10targets (which is a more common usage).

Caveat: when a target is going to be very long (like 20 tokens) with 10k targets, the resulting array will be 20 x 10k for ids, that can pile up quite fast memory usage. In that context, it could be much slower to pass everything at once. We need to benchmark that too.

We don't have benchmarking tests right now, but if a PR goes in I think a test should demonstrate the usage and have a clear comment at leat about this specific decision.
An example code can be found [in this colab example](https://colab.research.google.com/gist/EtaoinWu/0cf5b37882bd18bcc554d3da717a3974/fillmaskpipeline-test.ipynb). On the default Google machine that I wrote this notebook on, the version with a `targets` argument slows down significantly (100ish ms to 600ish ms).

> Caveat: when a target is going to be very long (like 20 tokens) with 10k targets, the resulting array will be 20 x 10k for ids, that can pile up quite fast memory usage. In that context, it could be much slower to pass everything at once. We need to benchmark that too.

The current behavior of `FillMaskPipeline` is that when a multi-token string is passed, only the first token is used. I doubt anyone would actually need this, because if someone want to choose a token from a subset of the vocabulary to fill into a mask, they usually know the subset exactly. Deliberately passing multi-token strings into `FillMaskPipeline` (and expecting it to tokenize them and drop all-but-first tokens) does not make much sense.

### Another discovery

When coding my example, I just discovered the bottleneck of the performance problem. When provided with a `targets` argument, `FillMaskPipeline` ignores its `top_k` parameter, which means that it has to output a whole list proportional to `len(targets)`, and that's the bottleneck (at least in my test). The code example above actually respects `top_k` parameter when a `targets` is present, hence much faster when constructing the return value. After the optimization, the code costs 200ish milliseconds.
> An example code can be found [in this colab example](https://colab.research.google.com/gist/EtaoinWu/0cf5b37882bd18bcc554d3da717a3974/fillmaskpipeline-test.ipynb). On the default Google machine that I wrote this notebook on, the version with a `targets` argument slows down significantly (100ish ms to 600ish ms).

Thanks, this will help ! 
> 
> > Caveat: when a target is going to be very long (like 20 tokens) with 10k targets, the resulting array will be 20 x 10k for ids, that can pile up quite fast memory usage. In that context, it could be much slower to pass everything at once. We need to benchmark that too.
> 
> The current behavior of `FillMaskPipeline` is that when a multi-token string is passed, only the first token is used. I doubt anyone would actually need this, because if someone want to choose a token from a subset of the vocabulary to fill into a mask, they usually know the subset exactly. Deliberately passing multi-token strings into `FillMaskPipeline` (and expecting it to tokenize them and drop all-but-first tokens) does not make much sense.

As a maintainer of a live product, I can tell you not everyone is aware of what happens behind a pipeline (and it is exactly why they exist, so we can abstract away all nitty gritty details of transformers). So it will happen that some users will try out those examples and be surprised at slowness. 
It's something that `pipelines` should try to address if possible. 

> ### Another discovery
> 
> When coding my example, I just discovered the bottleneck of the performance problem. When provided with a `targets` argument, `FillMaskPipeline` ignores its `top_k` parameter, which means that it has to output a whole list proportional to `len(targets)`, and that's the bottleneck (at least in my test). The code example above actually respects `top_k` parameter when a `targets` is present, hence much faster when constructing the return value. After the optimization, the code costs 200ish milliseconds.

Ok, I think I remember the `targets` being added and the decision was that if `top_k` > `len(targets)` we were not obliged of honoring `top_k` because it wouldn't make any sense. `top_k` < `len(targets)` should be honored though.



Resolving patch:
diff --git a/src/transformers/pipelines/fill_mask.py b/src/transformers/pipelines/fill_mask.py
--- a/src/transformers/pipelines/fill_mask.py
+++ b/src/transformers/pipelines/fill_mask.py
@@ -98,9 +98,9 @@ def __call__(self, *args, targets=None, top_k: Optional[int] = None, **kwargs):
             args (:obj:`str` or :obj:`List[str]`):
                 One or several texts (or one list of prompts) with masked tokens.
             targets (:obj:`str` or :obj:`List[str]`, `optional`):
-                When passed, the model will return the scores for the passed token or tokens rather than the top k
-                predictions in the entire vocabulary. If the provided targets are not in the model vocab, they will be
-                tokenized and the first resulting token will be used (with a warning).
+                When passed, the model will limit the scores to the passed targets instead of looking up in the whole
+                vocab. If the provided targets are not in the model vocab, they will be tokenized and the first
+                resulting token will be used (with a warning, and that might be slower).
             top_k (:obj:`int`, `optional`):
                 When passed, overrides the number of predictions to return.
 
@@ -115,25 +115,56 @@ def __call__(self, *args, targets=None, top_k: Optional[int] = None, **kwargs):
         inputs = self._parse_and_tokenize(*args, **kwargs)
         outputs = self._forward(inputs, return_tensors=True)
 
+        # top_k must be defined
+        if top_k is None:
+            top_k = self.top_k
+
         results = []
         batch_size = outputs.shape[0] if self.framework == "tf" else outputs.size(0)
 
         if targets is not None:
-            if len(targets) == 0 or len(targets[0]) == 0:
-                raise ValueError("At least one target must be provided when passed.")
             if isinstance(targets, str):
                 targets = [targets]
 
-            targets_proc = []
+            try:
+                vocab = self.tokenizer.get_vocab()
+            except Exception:
+                vocab = {}
+            target_ids = []
             for target in targets:
-                target_enc = self.tokenizer.tokenize(target)
-                if len(target_enc) > 1 or target_enc[0] == self.tokenizer.unk_token:
+                id_ = vocab.get(target, None)
+                if id_ is None:
+                    input_ids = self.tokenizer(
+                        target,
+                        add_special_tokens=False,
+                        return_attention_mask=False,
+                        return_token_type_ids=False,
+                        max_length=1,
+                        truncation=True,
+                    )["input_ids"]
+                    if len(input_ids) == 0:
+                        logger.warning(
+                            f"The specified target token `{target}` does not exist in the model vocabulary. "
+                            f"We cannot replace it with anything meaningful, ignoring it"
+                        )
+                        continue
+                    id_ = input_ids[0]
+                    # XXX: If users encounter this pass
+                    # it becomes pretty slow, so let's make sure
+                    # The warning enables them to fix the input to
+                    # get faster performance.
                     logger.warning(
                         f"The specified target token `{target}` does not exist in the model vocabulary. "
-                        f"Replacing with `{target_enc[0]}`."
+                        f"Replacing with `{self.tokenizer.convert_ids_to_tokens(id_)}`."
                     )
-                targets_proc.append(target_enc[0])
-            target_inds = np.array(self.tokenizer.convert_tokens_to_ids(targets_proc))
+                target_ids.append(id_)
+            target_ids = list(set(target_ids))
+            if len(target_ids) == 0:
+                raise ValueError("At least one target must be provided when passed.")
+            target_ids = np.array(target_ids)
+            # Cap top_k if there are targets
+            if top_k > target_ids.shape[0]:
+                top_k = target_ids.shape[0]
 
         for i in range(batch_size):
             input_ids = inputs["input_ids"][i]
@@ -147,14 +178,11 @@ def __call__(self, *args, targets=None, top_k: Optional[int] = None, **kwargs):
 
                 logits = outputs[i, masked_index.item(), :]
                 probs = tf.nn.softmax(logits)
-                if targets is None:
-                    topk = tf.math.top_k(probs, k=top_k if top_k is not None else self.top_k)
-                    values, predictions = topk.values.numpy(), topk.indices.numpy()
-                else:
-                    values = tf.gather_nd(probs, tf.reshape(target_inds, (-1, 1)))
-                    sort_inds = tf.reverse(tf.argsort(values), [0])
-                    values = tf.gather_nd(values, tf.reshape(sort_inds, (-1, 1))).numpy()
-                    predictions = target_inds[sort_inds.numpy()]
+                if targets is not None:
+                    probs = tf.gather_nd(probs, tf.reshape(target_ids, (-1, 1)))
+
+                topk = tf.math.top_k(probs, k=top_k)
+                values, predictions = topk.values.numpy(), topk.indices.numpy()
             else:
                 masked_index = torch.nonzero(input_ids == self.tokenizer.mask_token_id, as_tuple=False)
 
@@ -163,13 +191,11 @@ def __call__(self, *args, targets=None, top_k: Optional[int] = None, **kwargs):
 
                 logits = outputs[i, masked_index.item(), :]
                 probs = logits.softmax(dim=0)
-                if targets is None:
-                    values, predictions = probs.topk(top_k if top_k is not None else self.top_k)
-                else:
-                    values = probs[..., target_inds]
-                    sort_inds = list(reversed(values.argsort(dim=-1)))
-                    values = values[..., sort_inds]
-                    predictions = target_inds[sort_inds]
+
+                if targets is not None:
+                    probs = probs[..., target_ids]
+
+                values, predictions = probs.topk(top_k)
 
             for v, p in zip(values.tolist(), predictions.tolist()):
                 tokens = input_ids.numpy()


Test patch:
diff --git a/tests/test_pipelines_fill_mask.py b/tests/test_pipelines_fill_mask.py
--- a/tests/test_pipelines_fill_mask.py
+++ b/tests/test_pipelines_fill_mask.py
@@ -78,7 +78,8 @@ def test_torch_fill_mask(self):
     @require_torch
     def test_torch_fill_mask_with_targets(self):
         valid_inputs = ["My name is <mask>"]
-        valid_targets = [[" Teven", " Patrick", " Clara"], [" Sam"]]
+        # ' Sam' will yield a warning but work
+        valid_targets = [[" Teven", "ĠPatrick", "ĠClara"], ["ĠSam"], [" Sam"]]
         invalid_targets = [[], [""], ""]
         for model_name in self.small_models:
             unmasker = pipeline(task="fill-mask", model=model_name, tokenizer=model_name, framework="pt")
@@ -89,10 +90,34 @@ def test_torch_fill_mask_with_targets(self):
             for targets in invalid_targets:
                 self.assertRaises(ValueError, unmasker, valid_inputs, targets=targets)
 
+    @require_torch
+    def test_torch_fill_mask_with_targets_and_topk(self):
+        model_name = self.small_models[0]
+        unmasker = pipeline(task="fill-mask", model=model_name, tokenizer=model_name, framework="pt")
+        targets = [" Teven", "ĠPatrick", "ĠClara"]
+        top_k = 2
+        outputs = unmasker("My name is <mask>", targets=targets, top_k=top_k)
+
+        self.assertEqual(len(outputs), 2)
+
+    @require_torch
+    def test_torch_fill_mask_with_duplicate_targets_and_topk(self):
+        model_name = self.small_models[0]
+        unmasker = pipeline(task="fill-mask", model=model_name, tokenizer=model_name, framework="pt")
+        # String duplicates + id duplicates
+        targets = [" Teven", "ĠPatrick", "ĠClara", "ĠClara", " Clara"]
+        top_k = 10
+        outputs = unmasker("My name is <mask>", targets=targets, top_k=top_k)
+
+        # The target list contains duplicates, so we can't output more
+        # than them
+        self.assertEqual(len(outputs), 3)
+
     @require_tf
     def test_tf_fill_mask_with_targets(self):
         valid_inputs = ["My name is <mask>"]
-        valid_targets = [[" Teven", " Patrick", " Clara"], [" Sam"]]
+        # ' Sam' will yield a warning but work
+        valid_targets = [[" Teven", "ĠPatrick", "ĠClara"], ["ĠSam"], [" Sam"]]
         invalid_targets = [[], [""], ""]
         for model_name in self.small_models:
             unmasker = pipeline(task="fill-mask", model=model_name, tokenizer=model_name, framework="tf")
@@ -111,7 +136,7 @@ def test_torch_fill_mask_results(self):
             "My name is <mask>",
             "The largest city in France is <mask>",
         ]
-        valid_targets = [" Patrick", " Clara"]
+        valid_targets = ["ĠPatrick", "ĠClara"]
         for model_name in self.large_models:
             unmasker = pipeline(
                 task="fill-mask",
@@ -184,7 +209,7 @@ def test_tf_fill_mask_results(self):
             "My name is <mask>",
             "The largest city in France is <mask>",
         ]
-        valid_targets = [" Patrick", " Clara"]
+        valid_targets = ["ĠPatrick", "ĠClara"]
         for model_name in self.large_models:
             unmasker = pipeline(task="fill-mask", model=model_name, tokenizer=model_name, framework="tf", top_k=2)