Repository: keras-team/keras-nlp

Issue:
[keras-hub]Issue in running recently added models
When I am running the recently added models in keras-hub branch, I am getting an error.

```
import keras_nlp
import numpy as np
input_data = np.ones((2, 224, 224, 3), dtype="float32")

model = keras_nlp.models.ViTDetBackbone.from_preset("kaggle://keras/vitdet/keras/vitdet_base/2")
model(input_data)
```

```
Traceback (most recent call last):
  File "/Users/cosmo3769/projects/keras-nlp/test.py", line 64, in <module>
    model = keras_nlp.models.ViTDetBackbone.from_preset("kaggle://keras/vitdet/keras/vitdet_base/2")
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/cosmo3769/projects/keras-nlp/keras_nlp/src/models/backbone.py", line 212, in from_preset
    if not issubclass(preset_cls, cls):
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
TypeError: issubclass() arg 1 must be a class
```

**Note:** I am getting the same error in all the recently added models in keras-hub branch.

Hints from issue comments:
@cosmo3769: Ohh, I think kaggle models need to be updated for this to work.@Humbulani1234: I tried to investigate this issue further, and here is what I found out:

 The class function ```from_presets``` executes the line ```preset_cls = check_config_class(preset)``` which in turn executes the function ```check_config_class``` from ```keras_nlp.src.utils.preset_utils``` reproduced below:

```python
def check_config_class(
    preset,
    config_file=CONFIG_FILE,
):
    """Validate a preset is being loaded on the correct class."""
    config_path = get_file(preset, config_file)
    with open(config_path, encoding="utf-8") as config_file:
        config = json.load(config_file)
    return keras.saving.get_registered_object(config["registered_name"])
```
The above calls the function ```get_registered_object``` from ```keras.src.saving.object_registration``` and also reproduced below:

```python
def get_registered_object(name, custom_objects=None, module_objects=None):
    """Returns the class associated with `name` if it is registered with Keras.

    This function is part of the Keras serialization and deserialization
    framework. It maps strings to the objects associated with them for
    serialization/deserialization.

    Example:

    ```python
    def from_config(cls, config, custom_objects=None):
        if 'my_custom_object_name' in config:
            config['hidden_cls'] = tf.keras.saving.get_registered_object(
                config['my_custom_object_name'], custom_objects=custom_objects)
    ```
    Args:
        name: The name to look up.
        custom_objects: A dictionary of custom objects to look the name up in.
            Generally, custom_objects is provided by the user.
        module_objects: A dictionary of custom objects to look the name up in.
            Generally, module_objects is provided by midlevel library
            implementers.

    Returns:
        An instantiable class associated with `name`, or `None` if no such class
            exists.
    custom_objects_scope_dict = global_state.get_global_attribute(
        "custom_objects_scope_dict", {}
    )
    if name in custom_objects_scope_dict:
        return custom_objects_scope_dict[name]
    elif name in GLOBAL_CUSTOM_OBJECTS:
        return GLOBAL_CUSTOM_OBJECTS[name]
    elif custom_objects and name in custom_objects:
        return custom_objects[name]
    elif module_objects and name in module_objects:
        return module_objects[name]
    return None
```
which in turn calls ```get_global_attribute``` from ```keras.src.backend.common.global_state``, reproduced below`:

```python
def get_global_attribute(name, default=None, set_to_default=False):
    attr = getattr(GLOBAL_STATE_TRACKER, name, None)
    if attr is None and default is not None:
        attr = default
        if set_to_default:
            set_global_attribute(name, attr)
    return attr
```
I suspect this is where the problem is, the line ```attr = getattr(GLOBAL_STATE_TRACKER, name, None)``` when executed returns ```None``` object which is then assigned to ```attr``` and returned to the caller ```get_registered_object``` which in turn then returns ```None``` to ```check_config_class```. The ```preset_cls``` local variable will the be assigned the ```None``` object, hence when we execute the line ```if not issubclass(preset_cls, cls)``` from ```from_presets``` it raises the error ```TypeError: issubclass() arg 1 must be a class```.

I also noted that the downloaded ```config.json``` file has the ```registered name``` key given by: ```"keras_cv.models>ViTDetBackbone"```. I am not quite sure about whether the ```>``` should be a ```.```.

I hope this issue maybe be clarified further.

### The following is my detailed stacktrace produced by ```pdb```:

```
2024-08-27 12:41:16.184541: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
To enable the following instructions: AVX2 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
2024-08-27 12:41:17.025270: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
> /home/humbulani/keras-nlp/keras_nlp/src/models/backbone.py(214)from_preset()
-> preset_cls = check_config_class(preset)
(Pdb) s
--Call--
> /home/humbulani/keras-nlp/keras_nlp/src/utils/preset_utils.py(582)check_config_class()
-> def check_config_class(
(Pdb) s
> /home/humbulani/keras-nlp/keras_nlp/src/utils/preset_utils.py(587)check_config_class()
-> config_path = get_file(preset, config_file)
(Pdb) n
> /home/humbulani/keras-nlp/keras_nlp/src/utils/preset_utils.py(588)check_config_class()
-> with open(config_path, encoding="utf-8") as config_file:
(Pdb) n
> /home/humbulani/keras-nlp/keras_nlp/src/utils/preset_utils.py(589)check_config_class()
-> config = json.load(config_file)
(Pdb) n
> /home/humbulani/keras-nlp/keras_nlp/src/utils/preset_utils.py(588)check_config_class()
-> with open(config_path, encoding="utf-8") as config_file:
(Pdb) n
> /home/humbulani/keras-nlp/keras_nlp/src/utils/preset_utils.py(590)check_config_class()
-> return keras.saving.get_registered_object(config["registered_name"])
(Pdb) s
--Call--
> /home/humbulani/keras/env/lib/python3.10/site-packages/keras/src/saving/object_registration.py(188)get_registered_object()
-> @keras_export(
> /home/humbulani/keras/env/lib/python3.10/site-packages/keras/src/saving/object_registration.py(222)get_registered_object()
-> custom_objects_scope_dict = global_state.get_global_attribute(
(Pdb) s
> /home/humbulani/keras/env/lib/python3.10/site-packages/keras/src/saving/object_registration.py(223)get_registered_object()
-> "custom_objects_scope_dict", {}
(Pdb) s
> /home/humbulani/keras/env/lib/python3.10/site-packages/keras/src/saving/object_registration.py(222)get_registered_object()
-> custom_objects_scope_dict = global_state.get_global_attribute(
(Pdb) s
--Call--
> /home/humbulani/keras/env/lib/python3.10/site-packages/keras/src/backend/common/global_state.py(15)get_global_attribute()
-> def get_global_attribute(name, default=None, set_to_default=False):
(Pdb) list
> /home/humbulani/keras/env/lib/python3.10/site-packages/keras/src/backend/common/global_state.py(16)get_global_attribute()
-> attr = getattr(GLOBAL_STATE_TRACKER, name, None)
(Pdb) n
> /home/humbulani/keras/env/lib/python3.10/site-packages/keras/src/backend/common/global_state.py(17)get_global_attribute()
-> if attr is None and default is not None:
(Pdb) p attr
None
(Pdb) n
> /home/humbulani/keras/env/lib/python3.10/site-packages/keras/src/backend/common/global_state.py(18)get_global_attribute()
-> attr = default
(Pdb) s
> /home/humbulani/keras/env/lib/python3.10/site-packages/keras/src/backend/common/global_state.py(19)get_global_attribute()
-> if set_to_default:
(Pdb) s
> /home/humbulani/keras/env/lib/python3.10/site-packages/keras/src/backend/common/global_state.py(21)get_global_attribute()
-> return attr
(Pdb) s
--Return--
> /home/humbulani/keras/env/lib/python3.10/site-packages/keras/src/backend/common/global_state.py(21)get_global_attribute()->{}
-> return attr
(Pdb) s
> /home/humbulani/keras/env/lib/python3.10/site-packages/keras/src/saving/object_registration.py(225)get_registered_object()
-> if name in custom_objects_scope_dict:
(Pdb) p name
'keras_cv.models>ViTDetBackbone'
(Pdb) s
> /home/humbulani/keras/env/lib/python3.10/site-packages/keras/src/saving/object_registration.py(227)get_registered_object()
-> elif name in GLOBAL_CUSTOM_OBJECTS:
(Pdb) s
> /home/humbulani/keras/env/lib/python3.10/site-packages/keras/src/saving/object_registration.py(229)get_registered_object()
-> elif custom_objects and name in custom_objects:
(Pdb) s
> /home/humbulani/keras/env/lib/python3.10/site-packages/keras/src/saving/object_registration.py(231)get_registered_object()
-> elif module_objects and name in module_objects:
(Pdb) s
> /home/humbulani/keras/env/lib/python3.10/site-packages/keras/src/saving/object_registration.py(233)get_registered_object()
-> return None
(Pdb) s
--Return--
> /home/humbulani/keras/env/lib/python3.10/site-packages/keras/src/saving/object_registration.py(233)get_registered_object()->None
-> return None
(Pdb) s
--Return--
> /home/humbulani/keras-nlp/keras_nlp/src/utils/preset_utils.py(590)check_config_class()->None
-> return keras.saving.get_registered_object(config["registered_name"])
(Pdb) s
> /home/humbulani/keras-nlp/keras_nlp/src/models/backbone.py(215)from_preset()
-> if not issubclass(preset_cls, cls):
(Pdb) p preset_cls
None
(Pdb) s
TypeError: issubclass() arg 1 must be a class
> /home/humbulani/keras-nlp/keras_nlp/src/models/backbone.py(215)from_preset()
-> if not issubclass(preset_cls, cls):
(Pdb) s
--Return--
> /home/humbulani/keras-nlp/keras_nlp/src/models/backbone.py(215)from_preset()->None
-> if not issubclass(preset_cls, cls):
(Pdb) s
TypeError: issubclass() arg 1 must be a class
> /home/humbulani/keras-nlp/test.py(5)<module>()
-> model = keras_nlp.api.models.ViTDetBackbone.from_preset("kaggle://keras/vitdet/keras/vitdet_base/2")
(Pdb) s
--Return--
> /home/humbulani/keras-nlp/test.py(5)<module>()->None
-> model = keras_nlp.api.models.ViTDetBackbone.from_preset("kaggle://keras/vitdet/keras/vitdet_base/2")
(Pdb) s
Traceback (most recent call last):
  File "/home/humbulani/keras-nlp/test.py", line 5, in <module>
--Call--

(Pdb) n
--Call--
> /usr/lib/python3.10/codecs.py(319)decode()
-> def decode(self, input, final=False):
(Pdb) n

(Pdb) n
    model = keras_nlp.api.models.ViTDetBackbone.from_preset("kaggle://keras/vitdet/keras/vitdet_base/2")
  File "/home/humbulani/keras-nlp/keras_nlp/src/models/backbone.py", line 215, in from_preset

(Pdb) c
    if not issubclass(preset_cls, cls):
TypeError: issubclass() arg 1 must be a class
```@hsuyab: I think this error is coming up in one of the example codes as well. Here in the [getting started](https://keras.io/guides/keras_nlp/getting_started/) running `classifier = keras_nlp.models.BertClassifier.from_preset("bert_tiny_en_uncased_sst2")
` also gives the `TypeError: issubclass() arg 1 must be a class`.@mattdangerw: Thanks! See the error on `classifier = keras_nlp.models.BertClassifier.from_preset("bert_tiny_en_uncased_sst2")` taking a look.@mattdangerw: Ok! This should be now working as intende on master.

`classifier = keras_nlp.models.BertClassifier.from_preset("bert_tiny_en_uncased_sst2")` will work as expected.

`backbone = keras_nlp.models.ViTDetBackbone.from_preset("kaggle://keras/vitdet/keras/vitdet_base/2")` will thrown because we are attempting to revive a KerasCV model that is not available in the package. That is expected as well. But the error message will better.

For vit det (and other models we are adding to the hub) we will need to generate new kaggle assets before the model is ready for use. That's expected.

Solution patch:
diff --git a/keras_nlp/src/api_export.py b/keras_nlp/src/api_export.py
index 93e7b54c2f..9cf99eaff8 100644
--- a/keras_nlp/src/api_export.py
+++ b/keras_nlp/src/api_export.py
@@ -22,7 +22,16 @@
     namex = None
 
 
-def maybe_register_serializable(symbol):
+def maybe_register_serializable(path, symbol):
+    # If we have multiple export names, actually make sure to register these
+    # first. This makes sure we have a backward compat mapping of old serialized
+    # name to new class.
+    if isinstance(path, (list, tuple)):
+        for name in path:
+            name = name.split(".")[-1]
+            keras.saving.register_keras_serializable(
+                package="keras_nlp", name=name
+            )(symbol)
     if isinstance(symbol, types.FunctionType) or hasattr(symbol, "get_config"):
         # We register twice, first with the old name, second with the new name,
         # so loading still works under the old name.
@@ -39,7 +48,7 @@ def __init__(self, path):
             super().__init__(package="keras_nlp", path=path)
 
         def __call__(self, symbol):
-            maybe_register_serializable(symbol)
+            maybe_register_serializable(self.path, symbol)
             return super().__call__(symbol)
 
 else:
diff --git a/keras_nlp/src/utils/preset_utils.py b/keras_nlp/src/utils/preset_utils.py
index a935fd160b..83a7cd31d2 100644
--- a/keras_nlp/src/utils/preset_utils.py
+++ b/keras_nlp/src/utils/preset_utils.py
@@ -569,7 +569,16 @@ def load_serialized_object(config, **kwargs):
 
 def check_config_class(config):
     """Validate a preset is being loaded on the correct class."""
-    return keras.saving.get_registered_object(config["registered_name"])
+    registered_name = config["registered_name"]
+    cls = keras.saving.get_registered_object(registered_name)
+    if cls is None:
+        raise ValueError(
+            f"Attempting to load class {registered_name} with "
+            "`from_preset()`, but there is no class registered with Keras "
+            f"for {registered_name}. Make sure to register any custom "
+            "classes with `register_keras_serializable()`."
+        )
+    return cls
 
 
 def jax_memory_cleanup(layer):


Test patch:
diff --git a/keras_nlp/src/models/bert/bert_text_classifier_test.py b/keras_nlp/src/models/bert/bert_text_classifier_test.py
index 44b9a3d846..fa4f4cadfc 100644
--- a/keras_nlp/src/models/bert/bert_text_classifier_test.py
+++ b/keras_nlp/src/models/bert/bert_text_classifier_test.py
@@ -67,6 +67,15 @@ def test_saved_model(self):
             input_data=self.input_data,
         )
 
+    @pytest.mark.large
+    def test_smallest_preset(self):
+        self.run_preset_test(
+            cls=BertTextClassifier,
+            preset="bert_tiny_en_uncased_sst2",
+            input_data=self.input_data,
+            expected_output_shape=(2, 2),
+        )
+
     @pytest.mark.extra_large
     def test_all_presets(self):
         for preset in BertTextClassifier.presets:
diff --git a/keras_nlp/src/utils/preset_utils_test.py b/keras_nlp/src/utils/preset_utils_test.py
index 728e5a6967..9f38112a79 100644
--- a/keras_nlp/src/utils/preset_utils_test.py
+++ b/keras_nlp/src/utils/preset_utils_test.py
@@ -12,8 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import json
 import os
 
+import keras
 import pytest
 from absl.testing import parameterized
 
@@ -38,6 +40,15 @@ def test_preset_errors(self):
         with self.assertRaisesRegex(ValueError, "Unknown preset identifier"):
             AlbertTextClassifier.from_preset("snaggle://bort/bort/bort")
 
+        backbone = BertBackbone.from_preset("bert_tiny_en_uncased")
+        preset_dir = self.get_temp_dir()
+        config = keras.utils.serialize_keras_object(backbone)
+        config["registered_name"] = "keras_nlp>BortBackbone"
+        with open(os.path.join(preset_dir, CONFIG_FILE), "w") as config_file:
+            config_file.write(json.dumps(config, indent=4))
+        with self.assertRaisesRegex(ValueError, "class keras_nlp>BortBackbone"):
+            BertBackbone.from_preset(preset_dir)
+
     def test_upload_empty_preset(self):
         temp_dir = self.get_temp_dir()
         empty_preset = os.path.join(temp_dir, "empty")