Try changing the name of the argument

Author: mattdangerw

keras_nlp/src/models/preprocessor.py

@@ -89,7 +89,7 @@ def presets(cls):
     def from_preset(
         cls,
         preset,
-        load_task=False,
+        load_task_extras=False,
         **kwargs,
     ):
         """Instantiate a `keras_nlp.models.Preprocessor` from a model preset.
@@ -113,7 +113,7 @@ def from_preset(
         Args:
             preset: string. A built in preset identifier, a Kaggle Models
                 handle, a Hugging Face handle, or a path to a local directory.
-            load_task: bool. If `True`, load the saved task preprocessing
+            load_task_extras: bool. If `True`, load the saved task preprocessing
                 configuration from a `preprocessing.json`. You might use this to
                 restore the sequence length a model was fine-tuned with.
 
@@ -142,7 +142,7 @@ def from_preset(
         # Detect the correct subclass if we need to.
         if cls.backbone_cls != backbone_cls:
             cls = find_subclass(preset, cls, backbone_cls)
-        return loader.load_preprocessor(cls, load_task, **kwargs)
+        return loader.load_preprocessor(cls, load_task_extras, **kwargs)
 
     def save_to_preset(self, preset_dir):
         """Save preprocessor to a preset directory.

keras_nlp/src/models/preprocessor_test.py

@@ -115,7 +115,7 @@ def test_save_to_preset(self, cls, preset_name):
         self.assertEqual(set(tokenizer.file_assets), set(expected_assets))
 
         # Check restore.
-        restored = cls.from_preset(save_dir, load_task=True)
+        restored = cls.from_preset(save_dir, load_task_extras=True)
         self.assertEqual(preprocessor.get_config(), restored.get_config())
-        restored = cls.from_preset(save_dir, load_task=False)
+        restored = cls.from_preset(save_dir, load_task_extras=False)
         self.assertNotEqual(preprocessor.get_config(), restored.get_config())

keras_nlp/src/models/task.py

@@ -146,7 +146,7 @@ def from_preset(
         cls,
         preset,
         load_weights=True,
-        load_task=False,
+        load_task_extras=False,
         **kwargs,
     ):
         """Instantiate a `keras_nlp.models.Task` from a model preset.
@@ -175,7 +175,7 @@ def from_preset(
             load_weights: bool. If `True`, the backbone weights will be loaded
                 into the model architecture. If `False`, the weights will be
                 randomly initialized.
-            load_task: bool. If `True`, load the saved task configuration
+            load_task_extras: bool. If `True`, load the saved task configuration
                 from a `task.json` and any task specific weights from
                 `task.weights`. You might use this to load a classification
                 head for a model that has been saved with it.
@@ -206,7 +206,7 @@ def from_preset(
         # Detect the correct subclass if we need to.
         if cls.backbone_cls != backbone_cls:
             cls = find_subclass(preset, cls, backbone_cls)
-        return loader.load_task(cls, load_weights, load_task, **kwargs)
+        return loader.load_task(cls, load_weights, load_task_extras, **kwargs)
 
     def load_task_weights(self, filepath):
         """Load only the tasks specific weights not in the backbone."""

keras_nlp/src/models/task_test.py

@@ -132,7 +132,7 @@ def test_save_to_preset(self):
         self.assertEqual(BertTextClassifier, check_config_class(task_config))
 
         # Try loading the model from preset directory.
-        restored_task = TextClassifier.from_preset(save_dir, load_task=True)
+        restored_task = TextClassifier.from_preset(save_dir, load_task_extras=True)
 
         # Check the model output.
         data = ["the quick brown fox.", "the slow brown fox."]
@@ -142,7 +142,7 @@ def test_save_to_preset(self):
 
         # Load without head weights.
         restored_task = TextClassifier.from_preset(
-            save_dir, load_task=False, num_classes=2
+            save_dir, load_task_extras=False, num_classes=2
         )
         data = ["the quick brown fox.", "the slow brown fox."]
         # Full output unequal.

keras_nlp/src/utils/preset_utils.py

@@ -656,7 +656,7 @@ def load_tokenizer(self, cls, **kwargs):
         """Load a tokenizer layer from the preset."""
         raise NotImplementedError
 
-    def load_task(self, cls, load_weights, load_task, **kwargs):
+    def load_task(self, cls, load_weights, load_task_extras, **kwargs):
         """Load a task model from the preset.
 
         By default, we create a task from a backbone and preprocessor with
@@ -672,11 +672,11 @@ def load_task(self, cls, load_weights, load_task, **kwargs):
         if "preprocessor" not in kwargs:
             kwargs["preprocessor"] = self.load_preprocessor(
                 cls.preprocessor_cls,
-                load_task=load_task,
+                load_task_extras=load_task_extras,
             )
         return cls(**kwargs)
 
-    def load_preprocessor(self, cls, load_task, **kwargs):
+    def load_preprocessor(self, cls, load_task_extras, **kwargs):
         """Load a prepocessor layer from the preset.
 
         By default, we create a preprocessor from a tokenizer with default
@@ -705,23 +705,24 @@ def load_tokenizer(self, cls, **kwargs):
         tokenizer.load_preset_assets(self.preset)
         return tokenizer
 
-    def load_task(self, cls, load_weights, load_task, **kwargs):
+    def load_task(self, cls, load_weights, load_task_extras, **kwargs):
         # If there is no `task.json` or it's for the wrong class delegate to the
         # super class loader.
-        if not load_task:
-            return super().load_task(cls, load_weights, load_task, **kwargs)
+        if not load_task_extras:
+            return super().load_task(cls, load_weights, load_task_extras, **kwargs)
         if not check_file_exists(self.preset, TASK_CONFIG_FILE):
             raise ValueError(
                 "Saved preset has no `task.json`, cannot load the task config "
-                "from a file. Call `from_preset()` with `load_task=False` to "
-                "load the task from a backbone with library defaults."
+                "from a file. Call `from_preset()` with "
+                "`load_task_extras=False` to load the task from a backbone "
+                "with library defaults."
             )
         task_config = load_json(self.preset, TASK_CONFIG_FILE)
         if not issubclass(check_config_class(task_config), cls):
             raise ValueError(
                 f"Saved `task.json`does not match calling cls {cls}. Call "
-                "`from_preset()` with `load_task=False` to load the task from "
-                "a backbone with library defaults."
+                "`from_preset()` with `load_task_extras=False` to load the "
+                "task from a backbone with library defaults."
             )
         # We found a `task.json` with a complete config for our class.
         task = load_serialized_object(task_config, **kwargs)
@@ -736,21 +737,21 @@ def load_task(self, cls, load_weights, load_task, **kwargs):
             task.backbone.load_weights(backbone_weights)
         return task
 
-    def load_preprocessor(self, cls, load_task, **kwargs):
-        if not load_task:
-            return super().load_preprocessor(cls, load_task, **kwargs)
+    def load_preprocessor(self, cls, load_task_extras, **kwargs):
+        if not load_task_extras:
+            return super().load_preprocessor(cls, load_task_extras, **kwargs)
         if not check_file_exists(self.preset, PREPROCESSOR_CONFIG_FILE):
             raise ValueError(
                 "Saved preset has no `preprocessor.json`, cannot load the task "
                 "preprocessing config from a file. Call `from_preset()` with "
-                "`load_task=False` to load the preprocessor with library "
-                "defaults."
+                "`load_task_extras=False` to load the preprocessor with "
+                "library defaults."
             )
         preprocessor_json = load_json(self.preset, PREPROCESSOR_CONFIG_FILE)
         if not issubclass(check_config_class(preprocessor_json), cls):
             raise ValueError(
                 f"Saved `preprocessor.json`does not match calling cls {cls}. "
-                "Call `from_preset()` with `load_task=False` to "
+                "Call `from_preset()` with `load_task_extras=False` to "
                 "load the the preprocessor with library defaults."
             )
         # We found a `preprocessing.json` with a complete config for our class.