Repository: huggingface/transformers
Pull request: [Feature] Support `is_split_into_words` in the `TokenClassificationPipeline`.
# What does this PR do?
Fixes #30757

This PR solves the problem that when using an already tokenized input in the `TokenClassificationPipeline` by adding a `tokenizer_parameter` called `is_split_into_words`. I think we will need the aggregation code because some tokenizers are subword tokenizers, I have solved the problem that `start` and `end` was not correctly assigned when using the pipeline by reconstructing the sentence and create a map that stores the position of each word and pass it along with `words_id` to the `postprocess` and in `gather_pre_entities`, we use the map and `words_id` to produce the final position.
```py
from transformers import pipeline
from datasets import load_dataset

dataset = load_dataset("conll2003", split="test[:5]")

ner_pipeline = pipeline("token-classification", model="dslim/bert-base-NER", aggregation_strategy="first")

for example in dataset:
    tokens = example['tokens']

    predictions = ner_pipeline(tokens, is_split_into_words=True)

    for entity in predictions:
        print(entity)
```

Now this example doing inference on a batch of a preprocessed dataset. it gives this output.
```
[{'entity_group': 'MISC', 'score': np.float32(0.48014355), 'word': 'JAPAN', 'start': 9, 'end': 14}, {'entity_group': 'PER', 'score': np.float32(0.43156916), 'word': 'LUCKY', 'start': 19, 'end': 24}, {'entity_group': 'ORG', 'score': np.float32(0.64400136), 'word': 'CHINA', 'start': 31, 'end': 36}]
[{'entity_group': 'PER', 'score': np.float32(0.9979421), 'word': 'Nadim Ladki', 'start': 0, 'end': 11}]
[{'entity_group': 'LOC', 'score': np.float32(0.9851431), 'word': 'AL - AIN', 'start': 0, 'end': 6}, {'entity_group': 'LOC', 'score': np.float32(0.99941224), 'word': 'United Arab Emirates', 'start': 9, 'end': 29}]
[{'entity_group': 'LOC', 'score': np.float32(0.9998572), 'word': 'Japan', 'start': 0, 'end': 5}, {'entity_group': 'MISC', 'score': np.float32(0.99814206), 'word': 'Asian Cup', 'start': 33, 'end': 42}, {'entity_group': 'LOC', 'score': np.float32(0.9997607), 'word': 'Syria', 'start': 78, 'end': 83}, {'entity_group': 'MISC', 'score': np.float32(0.77703416), 'word': 'Group C', 'start': 89, 'end': 96}]
[{'entity_group': 'LOC', 'score': np.float32(0.9998673), 'word': 'China', 'start': 4, 'end': 9}, {'entity_group': 'LOC', 'score': np.float32(0.9997266), 'word': 'Uzbekistan', 'start': 119, 'end': 129}]
```
## Who can review?
@Rocketknight1 @ArthurZucker 


Problem statement:
TokenClassificationPipeline support is_split_into_words tokeniser parameter
### Feature request

The TokenClassificationPipeline currently sets a hardcoded tokeniser config within it sanitiser method. This prevents users from passing their own config to the tokeniser. 

It would be good to support some user input for tokeniser config. Especially for is_split_into_words as input data may be split already.

### Motivation

It is common for token classification datasets to be split into words already so that they match their labels. 

### Your contribution

I naivley anticipate this being a simple change, so I am happy to submit a PR for it. Though it would first be nice to see a discussion surrounding the feature and if it fits with the goals of Transformers.

Hints from issue discussion:
cc @ArthurZucker @Rocketknight1 
This makes sense to me, but I'm not super-familiar with that pipeline. I'd support a PR to allow some options to be passed through to the tokenizer, though, since that shouldn't have any backward compatibility issues!
Hello @swtb3, @amyeroberts, @Rocketknight1, @ArthurZucker,

I just saw this feature request because I've previously worked on this pipeline a few times ago. The `tokenizer_params` configuration hardcoded [here](https://github.com/huggingface/transformers/blob/f5c0fa9f6fe0eea2ad69bb1b03aff04824aa4870/src/transformers/pipelines/token_classification.py#L205-L209) was added in Transformers 4.28.0 to provide a chunkable mode for processing sequences longer than `tokenizer.model_max_length` which is only enabled when you specify the `stride` option to the pipeline.

Even though the tokenizer can accept `sentence`/`text` as a list of words while specifying `is_split_into_words`, here in the pipeline `sentence` should be a string currently and it is very critical as it is used to generate the offset mapping, then used to aggregate the entities. When using `return_offsets_mapping=True` with `is_split_into_words=True`, the generated offset mapping will be based on string in the list, not on the indexes of the list. For what I see, the easiest way to solve this would be to have a switch on the offset mapping: based on the string when `is_split_into_words=False` and based on the indexes of the list when `is_split_into_words=True`. 

Something like:

`is_split_into_words=False` -> `"HuggingFace is a company"` -> `{'entity_group': 'ORG', 'word': 'HuggingFace', 'start': 0, 'end': 11}` (with `aggregation_strategy` different from `None`)

`is_split_into_words=True` -> `['Hugging', 'Face', 'is', 'a', 'company']` -> `[{'entity': 'ORG', 'word': 'Hugging', 'start': 0, 'end': 1}, {'entity': 'ORG', 'word': 'Face', 'start': 1, 'end': 2}]` (`aggregation_strategy` should always be disabled here for simplicity. Even if some words can be splitted into multiple tokens, it's easy to pre-tokenize words with the corresponding tokenizer and then align labels as explained [here](https://huggingface.co/docs/transformers/en/tasks/token_classification#preprocess) with `tokenize_and_align_labels(examples)`)

If implemented, I strongly recommend to have `is_split_into_words` as a new added parameter than giving the user freedom to enter a raw `preprocess_params`/`tokenizer_params` which could potentially lead to misconfigured parameters and difficult for `_sanitize_parameters` to catch, leading to errors without being notified.

It is much more than just adding a parameter in a dict. I guess that the aim of this is to pre-annotate data while benefiting from the pipeline batching capability. I think that probably a lot of code will be left unused, especially in the `postprocess` method because aggregation is not really a problem as explained above.

Simply overriding `preprocess` and `postprocess` would be a better approach to have a relevant behavior and the ability to have a convenient format, so is it worth it to implement this?

I don't have the time to work on this currently but don't hesitate if you have any question

Commit message: removing unimportant lines

Patch to review:
diff --git a/src/transformers/pipelines/token_classification.py b/src/transformers/pipelines/token_classification.py
index e976ee9b1874..e69c02726a4b 100644
--- a/src/transformers/pipelines/token_classification.py
+++ b/src/transformers/pipelines/token_classification.py
@@ -1,6 +1,6 @@
 import types
 import warnings
-from typing import Any, Optional, Union, overload
+from typing import Any, Iterator, Optional, Union, overload
 
 import numpy as np
 
@@ -30,6 +30,8 @@ class TokenClassificationArgumentHandler(ArgumentHandler):
     """
 
     def __call__(self, inputs: Union[str, list[str]], **kwargs):
+        is_split_into_words = kwargs.get("is_split_into_words", False)
+
         if inputs is not None and isinstance(inputs, (list, tuple)) and len(inputs) > 0:
             inputs = list(inputs)
             batch_size = len(inputs)
@@ -37,7 +39,7 @@ def __call__(self, inputs: Union[str, list[str]], **kwargs):
             inputs = [inputs]
             batch_size = 1
         elif Dataset is not None and isinstance(inputs, Dataset) or isinstance(inputs, types.GeneratorType):
-            return inputs, None
+            return inputs, is_split_into_words, None
         else:
             raise ValueError("At least one input is required.")
 
@@ -47,7 +49,7 @@ def __call__(self, inputs: Union[str, list[str]], **kwargs):
                 offset_mapping = [offset_mapping]
             if len(offset_mapping) != batch_size:
                 raise ValueError("offset_mapping should have the same batch size as the input")
-        return inputs, offset_mapping
+        return inputs, is_split_into_words, offset_mapping
 
 
 class AggregationStrategy(ExplicitEnum):
@@ -135,6 +137,7 @@ class TokenClassificationPipeline(ChunkPipeline):
 
     def __init__(self, args_parser=TokenClassificationArgumentHandler(), *args, **kwargs):
         super().__init__(*args, **kwargs)
+
         self.check_model_type(
             TF_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING_NAMES
             if self.framework == "tf"
@@ -151,9 +154,11 @@ def _sanitize_parameters(
         ignore_subwords: Optional[bool] = None,
         aggregation_strategy: Optional[AggregationStrategy] = None,
         offset_mapping: Optional[list[tuple[int, int]]] = None,
+        is_split_into_words: Optional[bool] = False,
         stride: Optional[int] = None,
     ):
         preprocess_params = {}
+        preprocess_params["is_split_into_words"] = is_split_into_words
         if offset_mapping is not None:
             preprocess_params["offset_mapping"] = offset_mapping
 
@@ -230,8 +235,9 @@ def __call__(
         Classify each token of the text(s) given as inputs.
 
         Args:
-            inputs (`str` or `list[str]`):
-                One or several texts (or one list of texts) for token classification.
+            inputs (`str` or `List[str]`):
+                One or several texts (or one list of texts) for token classification. Can be pre-tokenized when
+                `is_split_into_words=True`.
 
         Return:
             A list or a list of list of `dict`: Each result comes as a list of dictionaries (one for each token in the
@@ -251,23 +257,62 @@ def __call__(
               exists if the offsets are available within the tokenizer
         """
 
-        _inputs, offset_mapping = self._args_parser(inputs, **kwargs)
+        _inputs, is_split_into_words, offset_mapping = self._args_parser(inputs, **kwargs)
+        kwargs["is_split_into_words"] = is_split_into_words
+        if is_split_into_words and not all(isinstance(input, list) for input in inputs):
+            return super().__call__([inputs], **kwargs)
         if offset_mapping:
             kwargs["offset_mapping"] = offset_mapping
 
         return super().__call__(inputs, **kwargs)
 
-    def preprocess(self, sentence, offset_mapping=None, **preprocess_params):
+    @overload
+    def preprocess(
+        self, sentence: str, is_split_into_words: bool = False, offset_mapping=None, **preprocess_params
+    ) -> Iterator[dict]: ...
+
+    @overload
+    def preprocess(
+        self, sentence: list[str], is_split_into_words: bool = True, offset_mapping=None, **preprocess_params
+    ) -> Iterator[dict]: ...
+
+    def preprocess(self, sentence, is_split_into_words, offset_mapping=None, **preprocess_params):
         tokenizer_params = preprocess_params.pop("tokenizer_params", {})
         truncation = True if self.tokenizer.model_max_length and self.tokenizer.model_max_length > 0 else False
+
+        word_to_chars_map = None
+        if is_split_into_words:
+            if not isinstance(sentence, list):
+                raise ValueError("When `is_split_into_words=True`, `sentence` must be a list of tokens.")
+            words = sentence
+            sentence = " ".join(words)  # Recreate the sentence string for later display and slicing
+            # This map will allows to convert back word => char indices
+            word_to_chars_map = []
+            char_offset = 0
+            for word in words:
+                word_to_chars_map.append((char_offset, char_offset + len(word)))
+                char_offset += len(word) + 1
+
+            # We use `words` as the actual input for the tokenizer
+            text_to_tokenize = words
+            tokenizer_params["is_split_into_words"] = True
+        else:
+            if not isinstance(sentence, str):
+                raise ValueError("When `is_split_into_words=False`, `sentence` must be an untokenized string.")
+            text_to_tokenize = sentence
+
         inputs = self.tokenizer(
-            sentence,
+            text_to_tokenize,
             return_tensors=self.framework,
             truncation=truncation,
             return_special_tokens_mask=True,
             return_offsets_mapping=self.tokenizer.is_fast,
             **tokenizer_params,
         )
+
+        if is_split_into_words and not self.tokenizer.is_fast:
+            raise ValueError("is_split_into_words=True is only supported with fast tokenizers.")
+
         inputs.pop("overflow_to_sample_mapping", None)
         num_chunks = len(inputs["input_ids"])
 
@@ -278,8 +323,12 @@ def preprocess(self, sentence, offset_mapping=None, **preprocess_params):
                 model_inputs = {k: v[i].unsqueeze(0) for k, v in inputs.items()}
             if offset_mapping is not None:
                 model_inputs["offset_mapping"] = offset_mapping
+
             model_inputs["sentence"] = sentence if i == 0 else None
             model_inputs["is_last"] = i == num_chunks - 1
+            if word_to_chars_map is not None:
+                model_inputs["word_ids"] = inputs.word_ids(i)
+                model_inputs["word_to_chars_map"] = word_to_chars_map
 
             yield model_inputs
 
@@ -289,6 +338,9 @@ def _forward(self, model_inputs):
         offset_mapping = model_inputs.pop("offset_mapping", None)
         sentence = model_inputs.pop("sentence")
         is_last = model_inputs.pop("is_last")
+        word_ids = model_inputs.pop("word_ids", None)
+        word_to_chars_map = model_inputs.pop("word_to_chars_map", None)
+
         if self.framework == "tf":
             logits = self.model(**model_inputs)[0]
         else:
@@ -301,6 +353,8 @@ def _forward(self, model_inputs):
             "offset_mapping": offset_mapping,
             "sentence": sentence,
             "is_last": is_last,
+            "word_ids": word_ids,
+            "word_to_chars_map": word_to_chars_map,
             **model_inputs,
         }
 
@@ -308,6 +362,10 @@ def postprocess(self, all_outputs, aggregation_strategy=AggregationStrategy.NONE
         if ignore_labels is None:
             ignore_labels = ["O"]
         all_entities = []
+
+        # Get map from the first output, it's the same for all chunks
+        word_to_chars_map = all_outputs[0].get("word_to_chars_map")
+
         for model_outputs in all_outputs:
             if self.framework == "pt" and model_outputs["logits"][0].dtype in (torch.bfloat16, torch.float16):
                 logits = model_outputs["logits"][0].to(torch.float32).numpy()
@@ -320,6 +378,7 @@ def postprocess(self, all_outputs, aggregation_strategy=AggregationStrategy.NONE
                 model_outputs["offset_mapping"][0] if model_outputs["offset_mapping"] is not None else None
             )
             special_tokens_mask = model_outputs["special_tokens_mask"][0].numpy()
+            word_ids = model_outputs.get("word_ids")
 
             maxes = np.max(logits, axis=-1, keepdims=True)
             shifted_exp = np.exp(logits - maxes)
@@ -330,7 +389,14 @@ def postprocess(self, all_outputs, aggregation_strategy=AggregationStrategy.NONE
                 offset_mapping = offset_mapping.numpy() if offset_mapping is not None else None
 
             pre_entities = self.gather_pre_entities(
-                sentence, input_ids, scores, offset_mapping, special_tokens_mask, aggregation_strategy
+                sentence,
+                input_ids,
+                scores,
+                offset_mapping,
+                special_tokens_mask,
+                aggregation_strategy,
+                word_ids=word_ids,
+                word_to_chars_map=word_to_chars_map,
             )
             grouped_entities = self.aggregate(pre_entities, aggregation_strategy)
             # Filter anything that is in self.ignore_labels
@@ -374,6 +440,8 @@ def gather_pre_entities(
         offset_mapping: Optional[list[tuple[int, int]]],
         special_tokens_mask: np.ndarray,
         aggregation_strategy: AggregationStrategy,
+        word_ids: Optional[list[Optional[int]]] = None,
+        word_to_chars_map: Optional[list[tuple[int, int]]] = None,
     ) -> list[dict]:
         """Fuse various numpy arrays into dicts with all the information needed for aggregation"""
         pre_entities = []
@@ -385,6 +453,15 @@ def gather_pre_entities(
             word = self.tokenizer.convert_ids_to_tokens(int(input_ids[idx]))
             if offset_mapping is not None:
                 start_ind, end_ind = offset_mapping[idx]
+
+                # If the input is pre-tokenized, we need to rescale the offsets to the absolute sentence.
+                if word_ids is not None and word_to_chars_map is not None:
+                    word_index = word_ids[idx]
+                    if word_index is not None:
+                        start_char, _ = word_to_chars_map[word_index]
+                        start_ind += start_char
+                        end_ind += start_char
+
                 if not isinstance(start_ind, int):
                     if self.framework == "pt":
                         start_ind = start_ind.item()
diff --git a/tests/pipelines/test_pipelines_token_classification.py b/tests/pipelines/test_pipelines_token_classification.py
index 5344ff980d4e..f14ec04d8e8f 100644
--- a/tests/pipelines/test_pipelines_token_classification.py
+++ b/tests/pipelines/test_pipelines_token_classification.py
@@ -308,6 +308,54 @@ def test_chunking(self):
             ],
         )
 
+    @require_torch
+    @slow
+    def test_is_split_into_words(self):
+        """
+        Tests the pipeline with pre-tokenized inputs (is_split_into_words=True)
+        and validates that the character offsets are correct.
+        """
+        token_classifier = pipeline(task="ner", model="dslim/bert-base-NER", aggregation_strategy="simple")
+
+        # Input is a list of words
+        words = ["Hello", "Sarah", "lives", "in", "New", "York"]
+
+        # The reconstructed sentence will be "Hello Sarah lives in New York"
+        # - "Sarah": starts at index 6, ends at 11
+        # - "New York": starts at index 21, ends at 29
+
+        output = token_classifier(words, is_split_into_words=True)
+
+        self.assertEqual(
+            nested_simplify(output),
+            [
+                {"entity_group": "PER", "score": ANY(float), "word": "Sarah", "start": 6, "end": 11},
+                {"entity_group": "LOC", "score": ANY(float), "word": "New York", "start": 21, "end": 29},
+            ],
+        )
+
+        # Also test batching with pre-tokenized inputs
+        words2 = ["My", "name", "is", "Wolfgang", "and", "I", "live", "in", "Berlin"]
+        batch_output = token_classifier([words, words2], is_split_into_words=True)
+
+        # Expected for second sentence ("My name is Wolfgang and I live in Berlin")
+        # - "Wolfgang": starts at 12, ends at 20
+        # - "Berlin": starts at 36, ends at 42
+
+        self.assertEqual(
+            nested_simplify(batch_output),
+            [
+                [
+                    {"entity_group": "PER", "score": ANY(float), "word": "Sarah", "start": 6, "end": 11},
+                    {"entity_group": "LOC", "score": ANY(float), "word": "New York", "start": 21, "end": 29},
+                ],
+                [
+                    {"entity_group": "PER", "score": ANY(float), "word": "Wolfgang", "start": 12, "end": 20},
+                    {"entity_group": "LOC", "score": ANY(float), "word": "Berlin", "start": 36, "end": 42},
+                ],
+            ],
+        )
+
     @require_torch
     def test_chunking_fast(self):
         # Note: We cannot run the test on "conflicts" on the chunking.
@@ -953,19 +1001,22 @@ def setUp(self):
     def test_simple(self):
         string = "This is a simple input"
 
-        inputs, offset_mapping = self.args_parser(string)
+        inputs, is_split_into_words, offset_mapping = self.args_parser(string)
         self.assertEqual(inputs, [string])
+        self.assertFalse(is_split_into_words)
         self.assertEqual(offset_mapping, None)
 
-        inputs, offset_mapping = self.args_parser([string, string])
+        inputs, is_split_into_words, offset_mapping = self.args_parser([string, string])
         self.assertEqual(inputs, [string, string])
+        self.assertFalse(is_split_into_words)
         self.assertEqual(offset_mapping, None)
 
-        inputs, offset_mapping = self.args_parser(string, offset_mapping=[(0, 1), (1, 2)])
+        inputs, is_split_into_words, offset_mapping = self.args_parser(string, offset_mapping=[(0, 1), (1, 2)])
         self.assertEqual(inputs, [string])
+        self.assertFalse(is_split_into_words)
         self.assertEqual(offset_mapping, [[(0, 1), (1, 2)]])
 
-        inputs, offset_mapping = self.args_parser(
+        inputs, is_split_into_words, offset_mapping = self.args_parser(
             [string, string], offset_mapping=[[(0, 1), (1, 2)], [(0, 2), (2, 3)]]
         )
         self.assertEqual(inputs, [string, string])


Code review comments:
diff_hunk: @@ -251,23 +257,62 @@ def __call__(
               exists if the offsets are available within the tokenizer
         """
 
-        _inputs, offset_mapping = self._args_parser(inputs, **kwargs)
+        _inputs, is_split_into_words, offset_mapping = self._args_parser(inputs, **kwargs)
+        kwargs["is_split_into_words"] = is_split_into_words
+        if is_split_into_words and not all(isinstance(input, list) for input in inputs):
+            return super().__call__([inputs], **kwargs)
         if offset_mapping:
             kwargs["offset_mapping"] = offset_mapping
 
         return super().__call__(inputs, **kwargs)
 
-    def preprocess(self, sentence, offset_mapping=None, **preprocess_params):
+    @overload
+    def preprocess(
+        self, sentence: str, is_split_into_words: bool = False, offset_mapping=None, **preprocess_params
+    ) -> Iterator[dict]: ...
+
+    @overload
+    def preprocess(
+        self, sentence: list[str], is_split_into_words: bool = True, offset_mapping=None, **preprocess_params
+    ) -> Iterator[dict]: ...
+
+    def preprocess(self, sentence, is_split_into_words, offset_mapping=None, **preprocess_params):
         tokenizer_params = preprocess_params.pop("tokenizer_params", {})
         truncation = True if self.tokenizer.model_max_length and self.tokenizer.model_max_length > 0 else False
+
+        word_to_chars_map = None
+        if is_split_into_words:
+            if not isinstance(sentence, list):
+                raise ValueError("When `is_split_into_words=True`, `sentence` must be a list of tokens.")
+            words = sentence
+            sentence = " ".join(words)  # Recreate the sentence string for later display and slicing
+            # This map will allows to convert back word => char indices
path: src/transformers/pipelines/token_classification.py
text: @user1:
Dumb question, but really `is_split_into_words` means that the input has been split into **tokens** rather than words, right? If any of the tokens aren't full words, this display will be wrong.

I'm not sure if that's fixable because different tokenizers use different ways of indicating word completions, I'm just not sure if this will always work perfectly!

@author:
Hey! Not a dumb question at all actually, a very insightful one.

You're right to question the assumptions around `is_split_into_words`, especially with the diversity in how tokenizers handle word boundaries. Luckily, this specific implementation accounts for that by checking whether a token is a subword using a simple but effective heuristic:
```py
is_subword = len(word) != len(word_ref)
```
This allows us to detect subword tokens even when tokenizers split words in unpredictable ways. The `aggregate_word` function then groups these subword tokens together correctly, leveraging the fact that tokenizers provide enough information to reconstruct the original word groupings.

That said, your point did make me realize there’s a deeper edge case here languages like Japanese or Chinese that don’t use whitespace as word boundaries.

So while the current code handles most whitespace-separated languages well, you're absolutely right that a more general, tokenizer-agnostic approach would be much trickier and probably need deeper integration with a smarter word alignment logic.

One idea I had for this is to expose a preprocessing parameter — something like a delimiter — so the user can explicitly define how their original words were separated. For whitespace-based languages, it could default to '  ', but for others, users could pass an empty string or custom logic. For example:

```py
sentence = delimiter.join(words)
```

Thanks for the question it helped me think more critically about generalization.

Let me know if you have any ideas or pointers on making this more robust I’d love to explore better ways to handle those cases.

diff_hunk: @@ -251,23 +257,62 @@ def __call__(
               exists if the offsets are available within the tokenizer
         """
 
-        _inputs, offset_mapping = self._args_parser(inputs, **kwargs)
+        _inputs, is_split_into_words, offset_mapping = self._args_parser(inputs, **kwargs)
+        kwargs["is_split_into_words"] = is_split_into_words
+        if is_split_into_words and not all(isinstance(input, list) for input in inputs):
+            return super().__call__([inputs], **kwargs)
         if offset_mapping:
             kwargs["offset_mapping"] = offset_mapping
 
         return super().__call__(inputs, **kwargs)
 
-    def preprocess(self, sentence, offset_mapping=None, **preprocess_params):
+    @overload
+    def preprocess(
+        self, sentence: str, is_split_into_words: bool = False, offset_mapping=None, **preprocess_params
+    ) -> Iterator[dict]: ...
+
+    @overload
+    def preprocess(
+        self, sentence: list[str], is_split_into_words: bool = True, offset_mapping=None, **preprocess_params
+    ) -> Iterator[dict]: ...
path: src/transformers/pipelines/token_classification.py
text: @user1:
The return type is the same in both cases so I'm not sure why we need the overload here!

@author:
You're totally right that overload was unnecessary. Not sure what I was thinking there 😅
The return type doesn’t change at all, so I’ve removed it from the code now.