[fix] special tokens moving

Author: 1kkiRen

TokenizerChanger/__init__.py

@@ -1,5 +1,5 @@
 """
-TokenizerChanger library v1.0.4
+TokenizerChanger library v1.0.5
 
 The Apache 2.0 License Copyright © Dmitrii Kuzmin
 """

TokenizerChanger/tokenizer_changer.py

@@ -89,12 +89,23 @@ def _move_special_tokens(self):
             self.state["added_tokens"][i]["id"] += (
                 len(self.state["model"]["vocab"]) - self.initial_length)
 
-        for i in range(len(self.state["post_processor"]["processors"])):
-            if 'special_tokens' in self.state["post_processor"]["processors"][i].keys():
-                for k in self.state["post_processor"]["processors"][i]["special_tokens"].keys():
-                    for j in tqdm(range(len(self.state["post_processor"]["processors"][i]["special_tokens"][k]['ids'])), desc="Moving special tokens"):
-                        self.state["post_processor"]["processors"][i]["special_tokens"][k]["ids"][j] += (
-                            len(self.state["model"]["vocab"]) - self.initial_length)
+        def process_special_tokens(obj):
+            if isinstance(obj, dict):
+                for key, value in obj.items():
+                    if key == "special_tokens" and isinstance(value, dict):
+                        for k in value.keys():
+                            if "ids" in value[k]:
+                                for j in tqdm(range(len(value[k]["ids"])), desc="Moving special tokens"):
+                                    value[k]["ids"][j] += (
+                                    len(self.state["model"]["vocab"]) - self.initial_length)
+                    else:
+                        process_special_tokens(value)
+
+            elif isinstance(obj, list):
+                for item in obj:
+                    process_special_tokens(item)
+
+        process_special_tokens(self.state.get("post_processor", {}))
 
     def _process_and_add_tokens(self, merge: list):
         processed_merge = ''.join(merge).replace(' ', '')
@@ -470,7 +481,7 @@ def updated_tokenizer(self):
         """
         self.__is_tokenizer()
 
-        if self.initial_length < len(self.state["model"]["vocab"]):
+        if self.initial_length != len(self.state["model"]["vocab"]):
             self._move_special_tokens()
 
         backend_tokenizer = Tokenizer.from_str(json.dumps(self.state))

docs/conf.py

@@ -14,7 +14,7 @@
 project = 'TokenizerChanger'
 copyright = '2024, 1kkiren'
 author = '1kkiren'
-release = '1.0.4'
+release = '1.0.5'
 
 # -- General configuration ---------------------------------------------------
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration