Add an option to disable default compilation (#1787)

keras_nlp/src/models/causal_lm.py

@@ -72,8 +72,6 @@ class CausalLM(Task):
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        # Default compilation.
-        self.compile()
 
     def compile(
         self,

keras_nlp/src/models/classifier.py

@@ -54,8 +54,6 @@ class Classifier(Task):
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        # Default compilation.
-        self.compile()
 
     def compile(
         self,

keras_nlp/src/models/masked_lm.py

@@ -45,8 +45,6 @@ class MaskedLM(Task):
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        # Default compilation.
-        self.compile()
 
     def compile(
         self,

keras_nlp/src/models/task.py

@@ -56,19 +56,27 @@ class Task(PipelineModel):
     to load a pre-trained config and weights. Calling `from_preset()` on a task
     will automatically instantiate a `keras_nlp.models.Backbone` and
     `keras_nlp.models.Preprocessor`.
+
+    Args:
+        compile: boolean, defaults to `True`. If `True` will compile the model
+            with default parameters on construction. Model can still be
+            recompiled with a new loss, optimizer and metrics before training.
     """
 
     backbone_cls = None
     preprocessor_cls = None
 
-    def __init__(self, *args, **kwargs):
+    def __init__(self, *args, compile=True, **kwargs):
         super().__init__(*args, **kwargs)
         self._functional_layer_ids = set(
             id(layer) for layer in self._flatten_layers()
         )
         self._initialized = True
         if self.backbone is not None:
             self.dtype_policy = self._backbone.dtype_policy
+        if compile:
+            # Default compilation.
+            self.compile()
 
     def preprocess_samples(self, x, y=None, sample_weight=None):
         if self.preprocessor is not None:

keras_nlp/src/tests/test_case.py

@@ -495,6 +495,10 @@ def run_task_test(
         task.preprocessor = None
         task.fit(ds.map(preprocessor))
         task.preprocessor = preprocessor
+        # Turn off default compilation, should error during `fit()`.
+        task = cls(**init_kwargs, compile=False)
+        with self.assertRaisesRegex(ValueError, "You must call `compile"):
+            task.fit(ds)
 
     def run_preset_test(
         self,