Update tokenization_indictrans.py

tokenization_indictrans.py

@@ -11,7 +11,10 @@ from transformers.tokenization_utils import PreTrainedTokenizer
 logger = logging.get_logger(__name__)
 
 SPIECE_UNDERLINE = "▁"
-SUPPORTED_LANGUAGES = [
+
+SPECIAL_TAGS = {
+    "_bt_",
+    "_ft_",
     "asm_Beng",
     "awa_Deva",
     "ben_Beng",
@@ -46,7 +49,7 @@ SUPPORTED_LANGUAGES = [
     "tel_Telu",
     "urd_Arab",
     "unr_Deva",
-]
+}
 
 VOCAB_FILES_NAMES = {
     "src_vocab_fp": "dict.SRC.json",
@@ -74,7 +77,7 @@ class IndicTransTokenizer(PreTrainedTokenizer):
         eos_token="</s>",
         pad_token="<pad>",
         do_lower_case=False,
-        **kwargs
+        **kwargs,
     ):
 
         self.src = True
@@ -124,7 +127,10 @@ class IndicTransTokenizer(PreTrainedTokenizer):
             pad_token=pad_token,
             **kwargs,
         )
-    
+
+    def add_new_special_tags(self, new_tags: List[str]):
+        SPECIAL_TAGS.update(new_tags)
+
     def _switch_to_input_mode(self):
         self.src = True
         self.padding_side = "left"
@@ -150,6 +156,16 @@ class IndicTransTokenizer(PreTrainedTokenizer):
         with open(path, "r", encoding="utf-8") as f:
             return json.load(f)
 
+    def _split_tags(self, tokens: List[str]) -> Tuple[List[str], List[str]]:
+        tags = [token for token in tokens if token in SPECIAL_TAGS]
+        tokens = [token for token in tokens if token not in SPECIAL_TAGS]
+        return tags, tokens
+
+    def _split_pads(self, tokens: List[str]) -> Tuple[List[str], List[str]]:
+        pads = [token for token in tokens if token == self.pad_token]
+        tokens = [token for token in tokens if token != self.pad_token]
+        return pads, tokens
+
     @property
     def src_vocab_size(self) -> int:
         return len(self.encoder)
@@ -183,27 +199,31 @@ class IndicTransTokenizer(PreTrainedTokenizer):
 
     def convert_tokens_to_string(self, tokens: List[str]) -> str:
         """Uses sentencepiece model for detokenization"""
-        pad_tokens = [token for token in tokens if token == self.pad_token]
-        tokens = [token for token in tokens if token != self.pad_token]
+        pads, tokens = self._split_pads(tokens)
+
         if self.src:
+
+            tags, non_tags = self._split_tags(tokens)
+
             return (
-                " ".join(pad_tokens)
+                " ".join(pads)
                 + " "
-                + " ".join(tokens[:2])
+                + " ".join(tags)
                 + " "
-                + "".join(tokens[2:]).replace(SPIECE_UNDERLINE, " ").strip()
+                + "".join(non_tags).replace(SPIECE_UNDERLINE, " ").strip()
             )
+
         return (
             "".join(tokens).replace(SPIECE_UNDERLINE, " ").strip()
             + " "
-            + " ".join(pad_tokens)
+            + " ".join(pads)
         )
 
     def _tokenize(self, text) -> List[str]:
         if self.src:
             tokens = text.split(" ")
-            tags = tokens[:2]
-            text = " ".join(tokens[2:])
+            tags, non_tags = self._split_tags(tokens)
+            text = " ".join(non_tags)
             tokens = self.current_spm.EncodeAsPieces(text)
             return tags + tokens
         else:
@@ -217,23 +237,25 @@ class IndicTransTokenizer(PreTrainedTokenizer):
         # We don't expect to process pairs, but leave the pair logic for API consistency
         return token_ids_0 + [self.eos_token_id] + token_ids_1 + [self.eos_token_id]
 
-    def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
+    def save_vocabulary(
+        self, save_directory: str, filename_prefix: Optional[str] = None
+    ) -> Tuple[str]:
         if not os.path.isdir(save_directory):
             logger.error(f"Vocabulary path ({save_directory}) should be a directory")
             return
-        
+
         src_spm_fp = os.path.join(save_directory, "model.SRC")
         tgt_spm_fp = os.path.join(save_directory, "model.TGT")
         src_vocab_fp = os.path.join(save_directory, "dict.SRC.json")
         tgt_vocab_fp = os.path.join(save_directory, "dict.TGT.json")
-        
+
         self._save_json(self.encoder, src_vocab_fp)
         self._save_json(self.decoder, tgt_vocab_fp)
-        
-        with open(src_spm_fp, 'wb') as f:
+
+        with open(src_spm_fp, "wb") as f:
             f.write(self.src_spm.serialized_model_proto())
-        
-        with open(tgt_spm_fp, 'wb') as f:
+
+        with open(tgt_spm_fp, "wb") as f:
             f.write(self.tgt_spm.serialized_model_proto())
 
-        return src_vocab_fp, tgt_vocab_fp, src_spm_fp, tgt_spm_fp
+        return src_vocab_fp, tgt_vocab_fp, src_spm_fp, tgt_spm_fp