def METHOD_NAME(
    raw_datasets: Union[hf_datasets.DatasetDict, hf_datasets.Dataset],
    model: torch.nn.Module,
    data_config: Union[Dict, attrdict.AttrDict],
    tokenizer: Any,
    text_column_name: str,
    label_column_name: str,
    label_to_id: Dict,
) -> Union[hf_datasets.Dataset, hf_datasets.DatasetDict]:
    padding = "max_length" if data_config.pad_to_max_length else False
    def tokenize_and_align_labels(
        examples,
    ):
        tokenized_inputs = tokenizer(
            examples[text_column_name],
            padding=padding,
            truncation=True,
            # We use this argument because the texts in our dataset are lists of words
            # (with a label for each word).
            is_split_into_words=True,
        )
        labels = []
        for i, label in enumerate(examples[label_column_name]):
            word_ids = tokenized_inputs.word_ids(batch_index=i)
            previous_word_idx = None
            label_ids = []
            for word_idx in word_ids:
                # Special tokens have a word id that is None. We set the label to -100 so they
                # are automatically ignored in the loss function.
                if word_idx is None:
                    label_ids.append(-100)
                # We set the label for the first token of each word.
                elif word_idx != previous_word_idx:
                    label_ids.append(label_to_id[label[word_idx]])
                # For the other tokens in a word, we set the label to either the current label
                # or -100, depending on the label_all_tokens flag.
                else:
                    label_ids.append(
                        label_to_id[label[word_idx]] if data_config.label_all_tokens else -100
                    )
                previous_word_idx = word_idx
            labels.append(label_ids)
        tokenized_inputs["labels"] = labels
        return tokenized_inputs
    tokenized_datasets = raw_datasets.map(
        tokenize_and_align_labels,
        num_proc=data_config.preprocessing_num_workers,
        load_from_cache_file=not data_config.overwrite_cache,
        batched=True,
    )
    for _, data in tokenized_datasets.items():
        hf.remove_unused_columns(model, data)
    return tokenized_datasets