Upload folder using huggingface_hub

README.md

@@ -49,6 +49,9 @@ Configure external evaluators with:
 The tokenizer accepts raw text through standard calls such as
 `tokenizer(text)`, `tokenizer(text, add_special_tokens=False)`, and
 `tokenizer(texts, padding=True, truncation=True, return_tensors="pt")`.
+It also accepts `return_offsets_mapping=True` for compatibility with
+completion-ranking evaluators that need suffix masks. The model supports
+`output_hidden_states=True` for representation extraction tasks.
 
 This export sets `patch_pathlib_utf8_open=true` in `config.json`. When loaded
 with `trust_remote_code=True`, the config installs a narrow Windows

hf/modeling_pinyin_code.py

@@ -139,10 +139,16 @@ class PinyinCodeModel(PinyinCodePreTrainedModel):
         attention_mask: torch.Tensor | None = None,
         inputs_embeds: torch.Tensor | None = None,
         position_ids: torch.Tensor | None = None,
+        output_hidden_states: bool | None = None,
         return_dict: bool | None = None,
         **kwargs,
     ) -> BaseModelOutput | tuple:
         return_dict = True if return_dict is None else return_dict
+        output_hidden_states = (
+            self.config.output_hidden_states
+            if output_hidden_states is None
+            else output_hidden_states
+        )
 
         if input_ids is None and inputs_embeds is None:
             raise ValueError("You must provide either input_ids or inputs_embeds")
@@ -173,14 +179,25 @@ class PinyinCodeModel(PinyinCodePreTrainedModel):
 
         x = inputs_embeds + self.position_embedding(position_ids)
         x = self.dropout(x)
+        all_hidden_states = (x,) if output_hidden_states else None
         for block in self.blocks:
             x = block(x, attention_mask=attention_mask)
+            if output_hidden_states:
+                all_hidden_states = all_hidden_states + (x,)
         hidden_states = self.ln_f(x)
+        if output_hidden_states:
+            all_hidden_states = all_hidden_states + (hidden_states,)
 
         if not return_dict:
-            return (hidden_states,)
+            output = (hidden_states,)
+            if output_hidden_states:
+                output = output + (all_hidden_states,)
+            return output
 
-        return BaseModelOutput(last_hidden_state=hidden_states)
+        return BaseModelOutput(
+            last_hidden_state=hidden_states,
+            hidden_states=all_hidden_states,
+        )
 
 
 class PinyinCodeForCausalLM(PinyinCodeModel, GenerationMixin):
@@ -230,9 +247,10 @@ class PinyinCodeForCausalLM(PinyinCodeModel, GenerationMixin):
         input_ids: torch.Tensor | None = None,
         attention_mask: torch.Tensor | None = None,
         labels: torch.Tensor | None = None,
-        inputs_embeds: torch.Tensor | None = None,
-        position_ids: torch.Tensor | None = None,
-        return_dict: bool | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        position_ids: torch.Tensor | None = None,
+        output_hidden_states: bool | None = None,
+        return_dict: bool | None = None,
         **kwargs,
     ) -> CausalLMOutput | tuple:
         return_dict = True if return_dict is None else return_dict
@@ -243,6 +261,7 @@ class PinyinCodeForCausalLM(PinyinCodeModel, GenerationMixin):
             attention_mask=attention_mask,
             inputs_embeds=inputs_embeds,
             position_ids=position_ids,
+            output_hidden_states=output_hidden_states,
             return_dict=True,
         )
         logits = self.lm_head(decoder_outputs.last_hidden_state)
@@ -255,8 +274,14 @@ class PinyinCodeForCausalLM(PinyinCodeModel, GenerationMixin):
                 ignore_index=-100,
             )
 
-        if not return_dict:
-            output = (logits,)
-            return ((loss,) + output) if loss is not None else output
-
-        return CausalLMOutput(loss=loss, logits=logits)
+        if not return_dict:
+            output = (logits,)
+            if decoder_outputs.hidden_states is not None:
+                output = output + (decoder_outputs.hidden_states,)
+            return ((loss,) + output) if loss is not None else output
+
+        return CausalLMOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=decoder_outputs.hidden_states,
+        )

hf/tokenization_pinyin_code.py

@@ -2,12 +2,13 @@
 
 from __future__ import annotations
 
-import logging
-import re
-import shutil
-import unicodedata
-from pathlib import Path
-from typing import Any
+import logging
+import math
+import re
+import shutil
+import unicodedata
+from pathlib import Path
+from typing import Any
 
 import sentencepiece as spm
 from transformers import PreTrainedTokenizer
@@ -77,7 +78,7 @@ def should_preserve_fallback_token(token: str) -> bool:
     return True
 
 
-class PinyinCodeTokenizer(PreTrainedTokenizer):
+class PinyinCodeTokenizer(PreTrainedTokenizer):
     """Slow tokenizer that preserves the existing SentencePiece model."""
 
     vocab_files_names = {"vocab_file": "tokenizer.model"}
@@ -276,30 +277,150 @@ class PinyinCodeTokenizer(PreTrainedTokenizer):
 
         return " ".join(tokens)
 
-    def _preprocess_tokenizer_input(self, value: Any) -> Any:
-        if value is None:
-            return None
-        if isinstance(value, str):
-            return self._preprocess_raw_text(value)
+    def _preprocess_tokenizer_input(self, value: Any) -> Any:
+        if value is None:
+            return None
+        if isinstance(value, str):
+            return self._preprocess_raw_text(value)
         if isinstance(value, tuple):
             return tuple(self._preprocess_tokenizer_input(item) for item in value)
         if isinstance(value, list):
-            return [self._preprocess_tokenizer_input(item) for item in value]
-        return value
-
-    def __call__(self, text=None, text_pair=None, *args, **kwargs):
-        if "text_target" in kwargs:
-            kwargs["text_target"] = self._preprocess_tokenizer_input(kwargs["text_target"])
-        if "text_pair_target" in kwargs:
-            kwargs["text_pair_target"] = self._preprocess_tokenizer_input(
-                kwargs["text_pair_target"]
+            return [self._preprocess_tokenizer_input(item) for item in value]
+        return value
+
+    def _non_content_token_ids(self) -> set[int]:
+        return {
+            token_id
+            for token_id in (
+                self.pad_token_id,
+                self.bos_token_id,
+                self.eos_token_id,
+                self.cls_token_id,
+                self.sep_token_id,
+                self.mask_token_id,
+            )
+            if token_id is not None
+        }
+
+    def _offset_source_text(self, value: Any, is_split_into_words: bool = False) -> str:
+        if value is None:
+            return ""
+        if isinstance(value, str):
+            return value
+        if isinstance(value, tuple):
+            return " ".join(self._offset_source_text(item) for item in value)
+        if isinstance(value, list):
+            separator = " " if is_split_into_words else ""
+            return separator.join(self._offset_source_text(item) for item in value)
+        return str(value)
+
+    def _synthetic_offset_mapping(self, text: Any, input_ids: Any, is_split_into_words: bool = False) -> list[tuple[int, int]]:
+        """Return slow-tokenizer-compatible offsets for evaluators that require them.
+
+        SentencePiece offsets are not available for this Python tokenizer because
+        raw Mandarin text is preprocessed into pinyin-code before encoding. These
+        spans conservatively distribute non-special tokens across the original
+        text so suffix/completion masking code can run without requiring a fast
+        tokenizer.
+        """
+        ids = input_ids.tolist() if hasattr(input_ids, "tolist") else list(input_ids)
+        source = self._offset_source_text(text, is_split_into_words=is_split_into_words)
+        source_length = len(source)
+        if not ids:
+            return []
+        if source_length == 0:
+            return [(0, 0) for _ in ids]
+
+        non_content_ids = self._non_content_token_ids()
+        content_positions = [
+            index for index, token_id in enumerate(ids) if int(token_id) not in non_content_ids
+        ]
+        if not content_positions:
+            return [(0, 0) for _ in ids]
+
+        offsets = [(0, 0) for _ in ids]
+        count = len(content_positions)
+        for ordinal, position in enumerate(content_positions):
+            start = math.floor(ordinal * source_length / count)
+            end = math.ceil((ordinal + 1) * source_length / count)
+            if end <= start:
+                end = min(source_length, start + 1)
+            offsets[position] = (start, end)
+        return offsets
+
+    def _with_optional_offsets(
+        self,
+        encoding,
+        original_text: Any,
+        return_offsets_mapping: bool,
+        is_split_into_words: bool = False,
+        return_tensors: str | None = None,
+    ):
+        if not return_offsets_mapping:
+            return encoding
+
+        input_ids = encoding["input_ids"]
+        tensor_input = hasattr(input_ids, "ndim")
+        input_ids_list = input_ids.tolist() if tensor_input else input_ids
+
+        is_batched = False
+        if tensor_input:
+            is_batched = input_ids.ndim > 1
+        elif input_ids_list and isinstance(input_ids_list[0], list):
+            is_batched = True
+
+        if is_batched:
+            if isinstance(original_text, list) and not is_split_into_words:
+                texts = original_text
+            else:
+                texts = [original_text] * len(input_ids_list)
+            offsets = [
+                self._synthetic_offset_mapping(text, ids, is_split_into_words=is_split_into_words)
+                for text, ids in zip(texts, input_ids_list)
+            ]
+        else:
+            offsets = self._synthetic_offset_mapping(
+                original_text,
+                input_ids_list,
+                is_split_into_words=is_split_into_words,
+            )
+
+        if return_tensors == "pt" or tensor_input:
+            try:
+                import torch
+
+                offsets = torch.tensor(offsets, dtype=torch.long)
+            except ImportError:
+                pass
+        encoding["offset_mapping"] = offsets
+        return encoding
+
+    def __call__(self, text=None, text_pair=None, *args, **kwargs):
+        original_text = text
+        return_offsets_mapping = bool(kwargs.pop("return_offsets_mapping", False))
+        is_split_into_words = bool(kwargs.get("is_split_into_words", False))
+        return_tensors = kwargs.get("return_tensors")
+
+        if "text_target" in kwargs:
+            kwargs["text_target"] = self._preprocess_tokenizer_input(kwargs["text_target"])
+        if "text_pair_target" in kwargs:
+            kwargs["text_pair_target"] = self._preprocess_tokenizer_input(
+                kwargs["text_pair_target"]
             )
-
-        text = self._preprocess_tokenizer_input(text)
-        text_pair = self._preprocess_tokenizer_input(text_pair)
-        if text_pair is None:
-            return super().__call__(text, *args, **kwargs)
-        return super().__call__(text, text_pair, *args, **kwargs)
+
+        text = self._preprocess_tokenizer_input(text)
+        text_pair = self._preprocess_tokenizer_input(text_pair)
+        if text_pair is None:
+            encoding = super().__call__(text, *args, **kwargs)
+        else:
+            encoding = super().__call__(text, text_pair, *args, **kwargs)
+        return self._with_optional_offsets(
+            encoding,
+            original_text,
+            return_offsets_mapping,
+            is_split_into_words=is_split_into_words,
+            return_tensors=return_tensors,
+        )
 
     def encode(self, text, text_pair=None, add_special_tokens=True, *args, **kwargs):
         kwargs["add_special_tokens"] = add_special_tokens
@@ -309,18 +430,43 @@ class PinyinCodeTokenizer(PreTrainedTokenizer):
             return super().encode(text, *args, **kwargs)
         return super().encode(text, text_pair, *args, **kwargs)
 
-    def encode_plus(self, text, text_pair=None, *args, **kwargs):
-        text = self._preprocess_tokenizer_input(text)
-        text_pair = self._preprocess_tokenizer_input(text_pair)
-        if text_pair is None:
-            return super().encode_plus(text, *args, **kwargs)
-        return super().encode_plus(text, text_pair, *args, **kwargs)
-
-    def batch_encode_plus(self, batch_text_or_text_pairs, *args, **kwargs):
-        batch_text_or_text_pairs = self._preprocess_tokenizer_input(
-            batch_text_or_text_pairs
-        )
-        return super().batch_encode_plus(batch_text_or_text_pairs, *args, **kwargs)
+    def encode_plus(self, text, text_pair=None, *args, **kwargs):
+        original_text = text
+        return_offsets_mapping = bool(kwargs.pop("return_offsets_mapping", False))
+        is_split_into_words = bool(kwargs.get("is_split_into_words", False))
+        return_tensors = kwargs.get("return_tensors")
+
+        text = self._preprocess_tokenizer_input(text)
+        text_pair = self._preprocess_tokenizer_input(text_pair)
+        if text_pair is None:
+            encoding = super().encode_plus(text, *args, **kwargs)
+        else:
+            encoding = super().encode_plus(text, text_pair, *args, **kwargs)
+        return self._with_optional_offsets(
+            encoding,
+            original_text,
+            return_offsets_mapping,
+            is_split_into_words=is_split_into_words,
+            return_tensors=return_tensors,
+        )
+
+    def batch_encode_plus(self, batch_text_or_text_pairs, *args, **kwargs):
+        original_batch = batch_text_or_text_pairs
+        return_offsets_mapping = bool(kwargs.pop("return_offsets_mapping", False))
+        is_split_into_words = bool(kwargs.get("is_split_into_words", False))
+        return_tensors = kwargs.get("return_tensors")
+
+        batch_text_or_text_pairs = self._preprocess_tokenizer_input(
+            batch_text_or_text_pairs
+        )
+        encoding = super().batch_encode_plus(batch_text_or_text_pairs, *args, **kwargs)
+        return self._with_optional_offsets(
+            encoding,
+            original_batch,
+            return_offsets_mapping,
+            is_split_into_words=is_split_into_words,
+            return_tensors=return_tensors,
+        )
 
     @property
     def vocab_size(self) -> int:

modeling_pinyin_code.py

@@ -139,10 +139,16 @@ class PinyinCodeModel(PinyinCodePreTrainedModel):
         attention_mask: torch.Tensor | None = None,
         inputs_embeds: torch.Tensor | None = None,
         position_ids: torch.Tensor | None = None,
+        output_hidden_states: bool | None = None,
         return_dict: bool | None = None,
         **kwargs,
     ) -> BaseModelOutput | tuple:
         return_dict = True if return_dict is None else return_dict
+        output_hidden_states = (
+            self.config.output_hidden_states
+            if output_hidden_states is None
+            else output_hidden_states
+        )
 
         if input_ids is None and inputs_embeds is None:
             raise ValueError("You must provide either input_ids or inputs_embeds")
@@ -173,14 +179,25 @@ class PinyinCodeModel(PinyinCodePreTrainedModel):
 
         x = inputs_embeds + self.position_embedding(position_ids)
         x = self.dropout(x)
+        all_hidden_states = (x,) if output_hidden_states else None
         for block in self.blocks:
             x = block(x, attention_mask=attention_mask)
+            if output_hidden_states:
+                all_hidden_states = all_hidden_states + (x,)
         hidden_states = self.ln_f(x)
+        if output_hidden_states:
+            all_hidden_states = all_hidden_states + (hidden_states,)
 
         if not return_dict:
-            return (hidden_states,)
+            output = (hidden_states,)
+            if output_hidden_states:
+                output = output + (all_hidden_states,)
+            return output
 
-        return BaseModelOutput(last_hidden_state=hidden_states)
+        return BaseModelOutput(
+            last_hidden_state=hidden_states,
+            hidden_states=all_hidden_states,
+        )
 
 
 class PinyinCodeForCausalLM(PinyinCodeModel, GenerationMixin):
@@ -230,9 +247,10 @@ class PinyinCodeForCausalLM(PinyinCodeModel, GenerationMixin):
         input_ids: torch.Tensor | None = None,
         attention_mask: torch.Tensor | None = None,
         labels: torch.Tensor | None = None,
-        inputs_embeds: torch.Tensor | None = None,
-        position_ids: torch.Tensor | None = None,
-        return_dict: bool | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        position_ids: torch.Tensor | None = None,
+        output_hidden_states: bool | None = None,
+        return_dict: bool | None = None,
         **kwargs,
     ) -> CausalLMOutput | tuple:
         return_dict = True if return_dict is None else return_dict
@@ -243,6 +261,7 @@ class PinyinCodeForCausalLM(PinyinCodeModel, GenerationMixin):
             attention_mask=attention_mask,
             inputs_embeds=inputs_embeds,
             position_ids=position_ids,
+            output_hidden_states=output_hidden_states,
             return_dict=True,
         )
         logits = self.lm_head(decoder_outputs.last_hidden_state)
@@ -255,8 +274,14 @@ class PinyinCodeForCausalLM(PinyinCodeModel, GenerationMixin):
                 ignore_index=-100,
             )
 
-        if not return_dict:
-            output = (logits,)
-            return ((loss,) + output) if loss is not None else output
-
-        return CausalLMOutput(loss=loss, logits=logits)
+        if not return_dict:
+            output = (logits,)
+            if decoder_outputs.hidden_states is not None:
+                output = output + (decoder_outputs.hidden_states,)
+            return ((loss,) + output) if loss is not None else output
+
+        return CausalLMOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=decoder_outputs.hidden_states,
+        )

tokenization_pinyin_code.py

@@ -2,12 +2,13 @@
 
 from __future__ import annotations
 
-import logging
-import re
-import shutil
-import unicodedata
-from pathlib import Path
-from typing import Any
+import logging
+import math
+import re
+import shutil
+import unicodedata
+from pathlib import Path
+from typing import Any
 
 import sentencepiece as spm
 from transformers import PreTrainedTokenizer
@@ -77,7 +78,7 @@ def should_preserve_fallback_token(token: str) -> bool:
     return True
 
 
-class PinyinCodeTokenizer(PreTrainedTokenizer):
+class PinyinCodeTokenizer(PreTrainedTokenizer):
     """Slow tokenizer that preserves the existing SentencePiece model."""
 
     vocab_files_names = {"vocab_file": "tokenizer.model"}
@@ -276,30 +277,150 @@ class PinyinCodeTokenizer(PreTrainedTokenizer):
 
         return " ".join(tokens)
 
-    def _preprocess_tokenizer_input(self, value: Any) -> Any:
-        if value is None:
-            return None
-        if isinstance(value, str):
-            return self._preprocess_raw_text(value)
+    def _preprocess_tokenizer_input(self, value: Any) -> Any:
+        if value is None:
+            return None
+        if isinstance(value, str):
+            return self._preprocess_raw_text(value)
         if isinstance(value, tuple):
             return tuple(self._preprocess_tokenizer_input(item) for item in value)
         if isinstance(value, list):
-            return [self._preprocess_tokenizer_input(item) for item in value]
-        return value
-
-    def __call__(self, text=None, text_pair=None, *args, **kwargs):
-        if "text_target" in kwargs:
-            kwargs["text_target"] = self._preprocess_tokenizer_input(kwargs["text_target"])
-        if "text_pair_target" in kwargs:
-            kwargs["text_pair_target"] = self._preprocess_tokenizer_input(
-                kwargs["text_pair_target"]
+            return [self._preprocess_tokenizer_input(item) for item in value]
+        return value
+
+    def _non_content_token_ids(self) -> set[int]:
+        return {
+            token_id
+            for token_id in (
+                self.pad_token_id,
+                self.bos_token_id,
+                self.eos_token_id,
+                self.cls_token_id,
+                self.sep_token_id,
+                self.mask_token_id,
+            )
+            if token_id is not None
+        }
+
+    def _offset_source_text(self, value: Any, is_split_into_words: bool = False) -> str:
+        if value is None:
+            return ""
+        if isinstance(value, str):
+            return value
+        if isinstance(value, tuple):
+            return " ".join(self._offset_source_text(item) for item in value)
+        if isinstance(value, list):
+            separator = " " if is_split_into_words else ""
+            return separator.join(self._offset_source_text(item) for item in value)
+        return str(value)
+
+    def _synthetic_offset_mapping(self, text: Any, input_ids: Any, is_split_into_words: bool = False) -> list[tuple[int, int]]:
+        """Return slow-tokenizer-compatible offsets for evaluators that require them.
+
+        SentencePiece offsets are not available for this Python tokenizer because
+        raw Mandarin text is preprocessed into pinyin-code before encoding. These
+        spans conservatively distribute non-special tokens across the original
+        text so suffix/completion masking code can run without requiring a fast
+        tokenizer.
+        """
+        ids = input_ids.tolist() if hasattr(input_ids, "tolist") else list(input_ids)
+        source = self._offset_source_text(text, is_split_into_words=is_split_into_words)
+        source_length = len(source)
+        if not ids:
+            return []
+        if source_length == 0:
+            return [(0, 0) for _ in ids]
+
+        non_content_ids = self._non_content_token_ids()
+        content_positions = [
+            index for index, token_id in enumerate(ids) if int(token_id) not in non_content_ids
+        ]
+        if not content_positions:
+            return [(0, 0) for _ in ids]
+
+        offsets = [(0, 0) for _ in ids]
+        count = len(content_positions)
+        for ordinal, position in enumerate(content_positions):
+            start = math.floor(ordinal * source_length / count)
+            end = math.ceil((ordinal + 1) * source_length / count)
+            if end <= start:
+                end = min(source_length, start + 1)
+            offsets[position] = (start, end)
+        return offsets
+
+    def _with_optional_offsets(
+        self,
+        encoding,
+        original_text: Any,
+        return_offsets_mapping: bool,
+        is_split_into_words: bool = False,
+        return_tensors: str | None = None,
+    ):
+        if not return_offsets_mapping:
+            return encoding
+
+        input_ids = encoding["input_ids"]
+        tensor_input = hasattr(input_ids, "ndim")
+        input_ids_list = input_ids.tolist() if tensor_input else input_ids
+
+        is_batched = False
+        if tensor_input:
+            is_batched = input_ids.ndim > 1
+        elif input_ids_list and isinstance(input_ids_list[0], list):
+            is_batched = True
+
+        if is_batched:
+            if isinstance(original_text, list) and not is_split_into_words:
+                texts = original_text
+            else:
+                texts = [original_text] * len(input_ids_list)
+            offsets = [
+                self._synthetic_offset_mapping(text, ids, is_split_into_words=is_split_into_words)
+                for text, ids in zip(texts, input_ids_list)
+            ]
+        else:
+            offsets = self._synthetic_offset_mapping(
+                original_text,
+                input_ids_list,
+                is_split_into_words=is_split_into_words,
+            )
+
+        if return_tensors == "pt" or tensor_input:
+            try:
+                import torch
+
+                offsets = torch.tensor(offsets, dtype=torch.long)
+            except ImportError:
+                pass
+        encoding["offset_mapping"] = offsets
+        return encoding
+
+    def __call__(self, text=None, text_pair=None, *args, **kwargs):
+        original_text = text
+        return_offsets_mapping = bool(kwargs.pop("return_offsets_mapping", False))
+        is_split_into_words = bool(kwargs.get("is_split_into_words", False))
+        return_tensors = kwargs.get("return_tensors")
+
+        if "text_target" in kwargs:
+            kwargs["text_target"] = self._preprocess_tokenizer_input(kwargs["text_target"])
+        if "text_pair_target" in kwargs:
+            kwargs["text_pair_target"] = self._preprocess_tokenizer_input(
+                kwargs["text_pair_target"]
             )
-
-        text = self._preprocess_tokenizer_input(text)
-        text_pair = self._preprocess_tokenizer_input(text_pair)
-        if text_pair is None:
-            return super().__call__(text, *args, **kwargs)
-        return super().__call__(text, text_pair, *args, **kwargs)
+
+        text = self._preprocess_tokenizer_input(text)
+        text_pair = self._preprocess_tokenizer_input(text_pair)
+        if text_pair is None:
+            encoding = super().__call__(text, *args, **kwargs)
+        else:
+            encoding = super().__call__(text, text_pair, *args, **kwargs)
+        return self._with_optional_offsets(
+            encoding,
+            original_text,
+            return_offsets_mapping,
+            is_split_into_words=is_split_into_words,
+            return_tensors=return_tensors,
+        )
 
     def encode(self, text, text_pair=None, add_special_tokens=True, *args, **kwargs):
         kwargs["add_special_tokens"] = add_special_tokens
@@ -309,18 +430,43 @@ class PinyinCodeTokenizer(PreTrainedTokenizer):
             return super().encode(text, *args, **kwargs)
         return super().encode(text, text_pair, *args, **kwargs)
 
-    def encode_plus(self, text, text_pair=None, *args, **kwargs):
-        text = self._preprocess_tokenizer_input(text)
-        text_pair = self._preprocess_tokenizer_input(text_pair)
-        if text_pair is None:
-            return super().encode_plus(text, *args, **kwargs)
-        return super().encode_plus(text, text_pair, *args, **kwargs)
-
-    def batch_encode_plus(self, batch_text_or_text_pairs, *args, **kwargs):
-        batch_text_or_text_pairs = self._preprocess_tokenizer_input(
-            batch_text_or_text_pairs
-        )
-        return super().batch_encode_plus(batch_text_or_text_pairs, *args, **kwargs)
+    def encode_plus(self, text, text_pair=None, *args, **kwargs):
+        original_text = text
+        return_offsets_mapping = bool(kwargs.pop("return_offsets_mapping", False))
+        is_split_into_words = bool(kwargs.get("is_split_into_words", False))
+        return_tensors = kwargs.get("return_tensors")
+
+        text = self._preprocess_tokenizer_input(text)
+        text_pair = self._preprocess_tokenizer_input(text_pair)
+        if text_pair is None:
+            encoding = super().encode_plus(text, *args, **kwargs)
+        else:
+            encoding = super().encode_plus(text, text_pair, *args, **kwargs)
+        return self._with_optional_offsets(
+            encoding,
+            original_text,
+            return_offsets_mapping,
+            is_split_into_words=is_split_into_words,
+            return_tensors=return_tensors,
+        )
+
+    def batch_encode_plus(self, batch_text_or_text_pairs, *args, **kwargs):
+        original_batch = batch_text_or_text_pairs
+        return_offsets_mapping = bool(kwargs.pop("return_offsets_mapping", False))
+        is_split_into_words = bool(kwargs.get("is_split_into_words", False))
+        return_tensors = kwargs.get("return_tensors")
+
+        batch_text_or_text_pairs = self._preprocess_tokenizer_input(
+            batch_text_or_text_pairs
+        )
+        encoding = super().batch_encode_plus(batch_text_or_text_pairs, *args, **kwargs)
+        return self._with_optional_offsets(
+            encoding,
+            original_batch,
+            return_offsets_mapping,
+            is_split_into_words=is_split_into_words,
+            return_tensors=return_tensors,
+        )
 
     @property
     def vocab_size(self) -> int: