Don't export model internals publicly (#1255)

Following our conventions, layers will be exposed standalone when they move from `models/xx/some_layer` -> `layers/some_layer`
keras-team · Sep 21, 2023 · bc21f21 · bc21f21
1 parent cecbeb2
commit bc21f21
Show file tree

Hide file tree

Showing 2 changed files with 0 additions and 4 deletions.
diff --git a/keras_nlp/models/xlnet/relative_attention.py b/keras_nlp/models/xlnet/relative_attention.py
@@ -15,7 +15,6 @@
 import math
 import string
 
-from keras_nlp.api_export import keras_nlp_export
 from keras_nlp.backend import keras
 from keras_nlp.backend import ops
 
@@ -76,7 +75,6 @@ def _rel_shift(x, klen=-1):
     return x
 
 
-@keras_nlp_export("keras_nlp.layers.TwoStreamRelativeAttention")
 class TwoStreamRelativeAttention(keras.layers.MultiHeadAttention):
     """Two-stream relative self-attention for XLNet.
 

diff --git a/keras_nlp/models/xlnet/xlnet_encoder.py b/keras_nlp/models/xlnet/xlnet_encoder.py
@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from keras_nlp.api_export import keras_nlp_export
 from keras_nlp.backend import keras
 from keras_nlp.backend import ops
 from keras_nlp.models.xlnet.relative_attention import TwoStreamRelativeAttention
@@ -22,7 +21,6 @@ def xlnet_kernel_initializer(stddev=0.02):
     return keras.initializers.TruncatedNormal(stddev=stddev)
 
 
-@keras_nlp_export("keras_nlp.layers.XLNetEncoder")
 class XLNetEncoder(keras.layers.Layer):
     """
     XLNet Encoder.