Added few more test cases in test_encode_decode_round_trip and modefi…

…ed the slow token (mask_token) to have AddedToken instance with lstrip=True
huggingface · Nov 30, 2023 · 99deea6 · 99deea6
1 parent acd4276
commit 99deea6
Show file tree

Hide file tree

Showing 2 changed files with 341 additions and 40 deletions.
diff --git a/src/transformers/models/rembert/tokenization_rembert.py b/src/transformers/models/rembert/tokenization_rembert.py
@@ -112,7 +112,7 @@ def __init__(
         **kwargs,
     ):
         # Mask token behave like a normal word, i.e. include the space before it
-        mask_token = AddedToken("[MASK]", lstrip=True, rstrip=False, normalized=False)
+        mask_token = AddedToken(mask_token, lstrip=True, rstrip=False) if isinstance(mask_token, str) else mask_token
 
         self.do_lower_case = do_lower_case
         self.remove_space = remove_space