Skip to content

Commit 9659340

Browse files
committed
Conditionally import tf text (#452)
In keeping with other layers, we should not rely on tf text being installed to import the library (this is useful for building keras.io for example).
1 parent 488b4ce commit 9659340

File tree

1 file changed

+10
-2
lines changed

1 file changed

+10
-2
lines changed

keras_nlp/tokenizers/byte_pair_tokenizer.py

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -24,10 +24,15 @@
2424
from typing import List
2525

2626
import tensorflow as tf
27-
import tensorflow_text as tf_text
2827
from tensorflow import keras
2928

3029
from keras_nlp.tokenizers import tokenizer
30+
from keras_nlp.utils.tf_utils import assert_tf_text_installed
31+
32+
try:
33+
import tensorflow_text as tf_text
34+
except ImportError:
35+
tf_text = None
3136

3237
# As python and TF handles special spaces differently, we need to
3338
# manually handle special spaces during string split.
@@ -220,6 +225,8 @@ def __init__(
220225
sequence_length=None,
221226
**kwargs,
222227
) -> None:
228+
assert_tf_text_installed(self.__class__.__name__)
229+
223230
# Check dtype and provide a default.
224231
if "dtype" not in kwargs or kwargs["dtype"] is None:
225232
kwargs["dtype"] = tf.int32
@@ -241,7 +248,8 @@ def __init__(
241248
else:
242249
raise ValueError(
243250
"Vocabulary must be an file path or dictionary mapping string "
244-
f"token to int ids. Received: `type(vocabulary)={type(vocabulary)}`."
251+
"token to int ids. Received: "
252+
f"`type(vocabulary)={type(vocabulary)}`."
245253
)
246254
if isinstance(merges, str):
247255
self.merges = [bp.rstrip() for bp in tf.io.gfile.GFile(merges)]

0 commit comments

Comments
 (0)