Add query_proj, value_proj to target names for enable_lora (keras-team#2107)

Author: abheesht17

keras_hub/src/models/backbone.py

@@ -186,14 +186,23 @@ def save_to_preset(self, preset_dir):
         saver = get_preset_saver(preset_dir)
         saver.save_backbone(self)
 
+    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.
+        """
+        return ["query_dense", "value_dense", "query", "value"]
+
     def enable_lora(self, rank):
         """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.
         """
-        target_names = ["query_dense", "value_dense", "query", "value"]
+        target_names = self.get_lora_target_names()
+
         self.trainable = True
         self._lora_enabled_layers = []
         self._lora_rank = rank

keras_hub/src/models/pali_gemma/pali_gemma_backbone.py

@@ -274,6 +274,13 @@ 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()
+
+        # Add these for `PaliGemmaVITAttention`.
+        target_names += ["query_proj", "value_proj"]
+        return target_names
+
     def get_config(self):
         config = super().get_config()
         config.update(