Update tokenizer test.

SamanehSaadat · Apr 22, 2024 · 33524af · 33524af
1 parent 2a3aca2
commit 33524af
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 9 deletions.
diff --git a/keras_nlp/models/preprocessor_test.py b/keras_nlp/models/preprocessor_test.py
@@ -91,7 +91,7 @@ def test_save_to_preset(self, cls, preset_name, tokenizer_type):
             vocab_filename = "vocabulary.txt"
             expected_assets = ["vocabulary.txt"]
 
-        # Check existence of files.
+        # Check existence of vocab file.
         vocab_path = os.path.join(
             save_dir, os.path.join(TOKENIZER_ASSET_DIR, vocab_filename)
         )

diff --git a/keras_nlp/tokenizers/tokenizer_test.py b/keras_nlp/tokenizers/tokenizer_test.py
@@ -28,7 +28,6 @@
 from keras_nlp.utils.preset_utils import TOKENIZER_ASSET_DIR
 from keras_nlp.utils.preset_utils import TOKENIZER_CONFIG_FILE
 from keras_nlp.utils.preset_utils import check_config_class
-from keras_nlp.utils.preset_utils import get_file
 
 
 class SimpleTokenizer(Tokenizer):
@@ -113,14 +112,11 @@ def test_save_to_preset(self, cls, preset_name, tokenizer_type):
             vocab_filename = "vocabulary.txt"
             expected_assets = ["vocabulary.txt"]
 
-        # Check existence of files.
-        self.assertTrue(
-            os.path.exists(
-                get_file(
-                    save_dir, os.path.join(TOKENIZER_ASSET_DIR, vocab_filename)
-                )
-            )
+        # Check existence of vocab file.
+        vocab_path = os.path.join(
+            save_dir, os.path.join(TOKENIZER_ASSET_DIR, vocab_filename)
         )
+        self.assertTrue(os.path.exists(vocab_path))
 
         # Check assets.
         self.assertEqual(set(tokenizer.file_assets), set(expected_assets))