Updated Large-scale multi-label text classification example for Keras v3

keras-team · Mar 5, 2024 · 2d76b3c · 2d76b3c
1 parent a3f7572
commit 2d76b3c
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/examples/nlp/multi_label_classification.py b/examples/nlp/multi_label_classification.py
@@ -30,8 +30,9 @@
 ## Imports
 """
 
-from tensorflow.keras import layers
-from tensorflow import keras
+import keras
+from keras import layers
+
 import tensorflow as tf
 
 from sklearn.model_selection import train_test_split
@@ -145,7 +146,7 @@
 """
 
 terms = tf.ragged.constant(train_df["terms"].values)
-lookup = tf.keras.layers.StringLookup(output_mode="multi_hot")
+lookup = layers.StringLookup(output_mode="multi_hot")
 lookup.adapt(terms)
 vocab = lookup.get_vocabulary()