Change set_vocabulary to set_proto

nkovela1 · Nov 21, 2023 · dc13ad0 · dc13ad0
1 parent 33f9bff
commit dc13ad0
Show file tree

Hide file tree

Showing 7 changed files with 16 additions and 17 deletions.
diff --git a/keras_nlp/models/albert/albert_tokenizer.py b/keras_nlp/models/albert/albert_tokenizer.py
@@ -94,8 +94,8 @@ def __init__(self, proto, **kwargs):
 
         super().__init__(proto=proto, **kwargs)
 
-    def set_vocabulary(self, proto):
-        super().set_vocabulary(proto)
+    def set_proto(self, proto):
+        super().set_proto(proto)
         if proto is not None:
             for token in [
                 self.cls_token,

diff --git a/keras_nlp/models/deberta_v3/deberta_v3_tokenizer.py b/keras_nlp/models/deberta_v3/deberta_v3_tokenizer.py
@@ -100,8 +100,8 @@ def __init__(self, proto, **kwargs):
 
         super().__init__(proto=proto, **kwargs)
 
-    def set_vocabulary(self, proto):
-        super().set_vocabulary(proto)
+    def set_proto(self, proto):
+        super().set_proto(proto)
         if proto is not None:
             for token in [self.cls_token, self.pad_token, self.sep_token]:
                 if token not in super().get_vocabulary():

diff --git a/keras_nlp/models/f_net/f_net_tokenizer.py b/keras_nlp/models/f_net/f_net_tokenizer.py
@@ -69,8 +69,8 @@ def __init__(self, proto, **kwargs):
         self.mask_token = "[MASK]"
         super().__init__(proto=proto, **kwargs)
 
-    def set_vocabulary(self, proto):
-        super().set_vocabulary(proto)
+    def set_proto(self, proto):
+        super().set_proto(proto)
         if proto is not None:
             for token in [
                 self.cls_token,

diff --git a/keras_nlp/models/t5/t5_tokenizer.py b/keras_nlp/models/t5/t5_tokenizer.py
@@ -78,10 +78,10 @@ def __init__(self, proto, **kwargs):
 
         super().__init__(proto=proto, **kwargs)
 
-    def set_vocabulary(self, proto):
-        super().set_vocabulary(proto)
+    def set_proto(self, proto):
+        super().set_proto(proto)
         if proto is not None:
-            for token in [self.pad_token]:
+            for token in [self.end_token, self.pad_token]:
                 if token not in self.get_vocabulary():
                     raise ValueError(
                         f"Cannot find token `'{token}'` in the provided "

diff --git a/keras_nlp/models/xlm_roberta/xlm_roberta_tokenizer.py b/keras_nlp/models/xlm_roberta/xlm_roberta_tokenizer.py
@@ -99,8 +99,8 @@ def __init__(self, proto, **kwargs):
 
         super().__init__(proto=proto, **kwargs)
 
-    def set_vocabulary(self, proto):
-        super().set_vocabulary(proto)
+    def set_proto(self, proto):
+        super().set_proto(proto)
         if proto is not None:
             self.mask_token_id = self.vocabulary_size() - 1
         else:

diff --git a/keras_nlp/tokenizers/sentence_piece_tokenizer.py b/keras_nlp/tokenizers/sentence_piece_tokenizer.py
@@ -128,19 +128,18 @@ def __init__(
 
         self.proto = None
         self.sequence_length = sequence_length
-        self.set_vocabulary(proto)
+        self.set_proto(proto)
 
     def save_assets(self, dir_path):
         path = os.path.join(dir_path, VOCAB_FILENAME)
         with open(path, "w") as file:
-            for token in self.proto:
-                file.write(f"{token}\n")
+            file.write(self.proto)
 
     def load_assets(self, dir_path):
         path = os.path.join(dir_path, VOCAB_FILENAME)
-        self.set_vocabulary(path)
+        self.set_proto(path)
 
-    def set_vocabulary(self, proto):
+    def set_proto(self, proto):
         if proto is None:
             self.proto = None
             self._sentence_piece = None

diff --git a/keras_nlp/tokenizers/sentence_piece_tokenizer_test.py b/keras_nlp/tokenizers/sentence_piece_tokenizer_test.py
@@ -161,7 +161,7 @@ def test_config(self):
         cloned_tokenizer = SentencePieceTokenizer.from_config(
             original_tokenizer.get_config()
         )
-        cloned_tokenizer.set_vocabulary(original_tokenizer.proto)
+        cloned_tokenizer.set_proto(original_tokenizer.proto)
         self.assertAllEqual(
             original_tokenizer(input_data),
             cloned_tokenizer(input_data),