fix get_lora_target_names function (#2167)

* fix target_names function

* fix in backbone

* update the get lora target names

* add setter

* address comment

* update docstring

Author: divyashreepathihalli

keras_hub/src/models/backbone.py

@@ -189,23 +189,26 @@ def save_to_preset(self, preset_dir, max_shard_size=10):
         saver = get_preset_saver(preset_dir)
         saver.save_backbone(self, max_shard_size=max_shard_size)
 
-    def get_lora_target_names(self):
-        """Returns list of layer names which are to be LoRA-fied.
-
-        Subclasses can override this method if the names of layers to be
-        LoRa-fied are different.
-        """
+    def default_lora_layer_names(self):
+        """Returns list of layer names which are to be LoRA-fied."""
         return ["query_dense", "value_dense", "query", "value"]
 
-    def enable_lora(self, rank, target_names=None):
+    def enable_lora(self, rank, target_layer_names=None):
         """Enable Lora on the backbone.
 
         Calling this method will freeze all weights on the backbone,
         while enabling Lora on the query & value `EinsumDense` layers
         of the attention layers.
+
+        Args:
+            rank: The rank of the LoRA factorization.
+            target_layer_names: A list of strings, the names of the layers to
+                apply LoRA to. If `None`, this will be populated with the
+                default LoRA layer names as returned by
+                `backbone.default_lora_layer_names()`.
         """
-        if target_names is None:
-            target_names = self.get_lora_target_names()
+        if target_layer_names is None:
+            target_layer_names = self.default_lora_layer_names()
         self.trainable = True
         self._lora_enabled_layers = []
         self._lora_rank = rank
@@ -214,7 +217,7 @@ def enable_lora(self, rank, target_names=None):
         all_layers = self._flatten_layers(include_self=False)
         all_layers = [lyr for lyr in all_layers if lyr.weights]
         for i, layer in enumerate(all_layers):
-            for name in target_names:
+            for name in target_layer_names:
                 if layer.name == name:
                     if hasattr(layer, "enable_lora"):
                         layer.trainable = True

keras_hub/src/models/gemma/gemma_lora_test.py

@@ -50,7 +50,7 @@ def test_lora_fine_tuning(self):
     def test_lora_fine_tuning_target_names(self):
         # Set up backbone and preprocessor.
         backbone = GemmaBackbone(**self._init_kwargs)
-        backbone.enable_lora(4, target_names=["query"])
+        backbone.enable_lora(4, target_layer_names=["query"])
         # 4 layers, 2 weights per layer
         self.assertLen(backbone.trainable_weights, 2 * 2)
         self.assertLen(backbone.non_trainable_weights, 20)

keras_hub/src/models/gemma3/gemma3_backbone.py

@@ -402,8 +402,8 @@ def get_config(self):
         )
         return config
 
-    def get_lora_target_names(self):
-        target_names = super().get_lora_target_names()
+    def default_lora_layer_names(self):
+        target_names = super().default_lora_layer_names()
 
         # Add these for `Gemma3VITAttention`.
         if not self.text_only_model:

keras_hub/src/models/pali_gemma/pali_gemma_backbone.py

@@ -274,8 +274,8 @@ def __init__(
         # Keep the image_sequence_length as a backbone property for easy access.
         self.image_sequence_length = self.vit_encoder.image_sequence_length
 
-    def get_lora_target_names(self):
-        target_names = super().get_lora_target_names()
+    def default_lora_layer_names(self):
+        target_names = super().default_lora_layer_names()
 
         # Add these for `PaliGemmaVITAttention`.
         target_names += ["query_proj", "value_proj"]