Move gpt model code into a directory

keras_nlp/models/__init__.py

@@ -20,11 +20,11 @@
 from keras_nlp.models.bert.bert_models import BertTiny
 from keras_nlp.models.bert.bert_preprocessing import BertPreprocessor
 from keras_nlp.models.bert.bert_tasks import BertClassifier
-from keras_nlp.models.gpt2 import Gpt2Base
-from keras_nlp.models.gpt2 import Gpt2Custom
-from keras_nlp.models.gpt2 import Gpt2ExtraLarge
-from keras_nlp.models.gpt2 import Gpt2Large
-from keras_nlp.models.gpt2 import Gpt2Medium
+from keras_nlp.models.gpt2.gpt2_models import Gpt2Base
+from keras_nlp.models.gpt2.gpt2_models import Gpt2Custom
+from keras_nlp.models.gpt2.gpt2_models import Gpt2ExtraLarge
+from keras_nlp.models.gpt2.gpt2_models import Gpt2Large
+from keras_nlp.models.gpt2.gpt2_models import Gpt2Medium
 from keras_nlp.models.roberta import RobertaBase
 from keras_nlp.models.roberta import RobertaClassifier
 from keras_nlp.models.roberta import RobertaCustom

keras_nlp/models/gpt2/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2022 The KerasNLP Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

keras_nlp/models/gpt2.py → keras_nlp/models/gpt2/gpt2_models.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""GPT-2 model configurable class, preconfigured versions, and task heads."""
+"""GPT-2 backbone models."""
 
 import tensorflow as tf
 from tensorflow import keras

keras_nlp/models/gpt2_test.py → keras_nlp/models/gpt2/gpt2_models_test.py

@@ -11,20 +11,21 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""Tests for GPT-2 model."""
+"""Test for GPT-2 backbone models."""
 
 import os
 
 import tensorflow as tf
 from absl.testing import parameterized
 from tensorflow import keras
 
-from keras_nlp.models import gpt2
+from keras_nlp.models.gpt2.gpt2_models import Gpt2Base
+from keras_nlp.models.gpt2.gpt2_models import Gpt2Custom
 
 
 class Gpt2Test(tf.test.TestCase, parameterized.TestCase):
     def setUp(self):
-        self.model = gpt2.Gpt2Custom(
+        self.model = Gpt2Custom(
             vocabulary_size=1000,
             num_layers=2,
             num_heads=2,
@@ -63,22 +64,22 @@ def test_variable_sequence_length_call_gpt2(self):
             self.model(input_data)
 
     def test_valid_call_gpt2_base(self):
-        model = gpt2.Gpt2Base(vocabulary_size=1000, name="gpt2_base_test")
+        model = Gpt2Base(vocabulary_size=1000, name="gpt2_base_test")
         model(self.input_batch)
 
     @parameterized.named_parameters(
         ("jit_compile_false", False), ("jit_compile_true", True)
     )
     def test_gpt2_base_compile(self, jit_compile):
-        model = gpt2.Gpt2Base(vocabulary_size=1000, name="gpt2_base_test")
+        model = Gpt2Base(vocabulary_size=1000, name="gpt2_base_test")
         model.compile(jit_compile=jit_compile)
         model.predict(self.input_batch)
 
     @parameterized.named_parameters(
         ("jit_compile_false", False), ("jit_compile_true", True)
     )
     def test_gpt2_base_compile_batched_ds(self, jit_compile):
-        model = gpt2.Gpt2Base(vocabulary_size=1000, name="gpt2_base_test")
+        model = Gpt2Base(vocabulary_size=1000, name="gpt2_base_test")
         model.compile(jit_compile=jit_compile)
         model.predict(self.input_dataset)