Add generic export_to_transformers to the base classes (#2346)

* export_to_transformers

* Address comments

* Fixes AutoTokenizer and AutoModel compatibility

* Add tests

* Add lm_head

* export_to_transformers

* Address comments

* Fixes AutoTokenizer and AutoModel compatibility

* Add tests

* Add lm_head

* resolve gemma_test

* Address comments and add llama model export support

* Move llama to a new PR

* address nits

Author: Bond099

keras_hub/src/models/backbone.py

@@ -277,3 +277,19 @@ def load_lora_weights(self, filepath):
             layer.lora_kernel_a.assign(lora_kernel_a)
             layer.lora_kernel_b.assign(lora_kernel_b)
         store.close()
+
+    def export_to_transformers(self, path):
+        """Export the backbone model to HuggingFace Transformers format.
+
+        This saves the backbone's configuration and weights in a format
+        compatible with HuggingFace Transformers. For unsupported model
+        architectures, a ValueError is raised.
+
+        Args:
+            path: str. Path to save the exported model.
+        """
+        from keras_hub.src.utils.transformers.export.hf_exporter import (
+            export_backbone,
+        )
+
+        export_backbone(self, path)

keras_hub/src/models/backbone_test.py

@@ -5,6 +5,7 @@
 
 from keras_hub.src.models.backbone import Backbone
 from keras_hub.src.models.bert.bert_backbone import BertBackbone
+from keras_hub.src.models.gemma.gemma_backbone import GemmaBackbone
 from keras_hub.src.models.gpt2.gpt2_backbone import GPT2Backbone
 from keras_hub.src.tests.test_case import TestCase
 from keras_hub.src.utils.preset_utils import CONFIG_FILE
@@ -105,3 +106,40 @@ def test_save_to_preset(self):
         ref_out = backbone(data)
         new_out = restored_backbone(data)
         self.assertAllClose(ref_out, new_out)
+
+    def test_export_supported_model(self):
+        backbone_config = {
+            "vocabulary_size": 1000,
+            "num_layers": 2,
+            "num_query_heads": 4,
+            "num_key_value_heads": 1,
+            "hidden_dim": 512,
+            "intermediate_dim": 1024,
+            "head_dim": 128,
+        }
+        backbone = GemmaBackbone(**backbone_config)
+        export_path = os.path.join(self.get_temp_dir(), "export_backbone")
+        backbone.export_to_transformers(export_path)
+        # Basic check: config file exists
+        self.assertTrue(
+            os.path.exists(os.path.join(export_path, "config.json"))
+        )
+
+    def test_export_unsupported_model(self):
+        backbone_config = {
+            "vocabulary_size": 1000,
+            "num_layers": 2,
+            "num_query_heads": 4,
+            "num_key_value_heads": 1,
+            "hidden_dim": 512,
+            "intermediate_dim": 1024,
+            "head_dim": 128,
+        }
+
+        class UnsupportedBackbone(GemmaBackbone):
+            pass
+
+        backbone = UnsupportedBackbone(**backbone_config)
+        export_path = os.path.join(self.get_temp_dir(), "unsupported")
+        with self.assertRaises(ValueError):
+            backbone.export_to_transformers(export_path)

keras_hub/src/models/causal_lm.py

@@ -392,3 +392,24 @@ def postprocess(x):
             outputs = [postprocess(x) for x in outputs]
 
         return self._normalize_generate_outputs(outputs, input_is_scalar)
+
+    def export_to_transformers(self, path):
+        """Export the full CausalLM model to HuggingFace Transformers format.
+
+        This exports the trainable model, tokenizer, and configurations in a
+        format compatible with HuggingFace Transformers. For unsupported model
+        architectures, a ValueError is raised.
+
+        If the preprocessor is attached (default), both the trainable model and
+        tokenizer are exported. To export only the trainable model, set
+        `self.preprocessor = None` before calling this method, then export the
+        preprocessor separately via `preprocessor.export_to_transformers(path)`.
+
+        Args:
+            path: str. Path to save the exported model.
+        """
+        from keras_hub.src.utils.transformers.export.hf_exporter import (
+            export_to_safetensors,
+        )
+
+        export_to_safetensors(self, path)

keras_hub/src/models/causal_lm_preprocessor.py

@@ -180,3 +180,17 @@ def sequence_length(self, value):
         self._sequence_length = value
         if self.packer is not None:
             self.packer.sequence_length = value
+
+    def export_to_transformers(self, path):
+        """Export the preprocessor to HuggingFace Transformers format.
+
+        Args:
+            path: str. Path to save the exported preprocessor/tokenizer.
+        """
+        if self.tokenizer is None:
+            raise ValueError("Preprocessor must have a tokenizer for export.")
+        from keras_hub.src.utils.transformers.export.hf_exporter import (
+            export_tokenizer,
+        )
+
+        export_tokenizer(self.tokenizer, path)

keras_hub/src/models/causal_lm_preprocessor_test.py

@@ -1,7 +1,13 @@
+import os
+
 import pytest
 
 from keras_hub.src.models.bert.bert_tokenizer import BertTokenizer
 from keras_hub.src.models.causal_lm_preprocessor import CausalLMPreprocessor
+from keras_hub.src.models.gemma.gemma_causal_lm_preprocessor import (
+    GemmaCausalLMPreprocessor,
+)
+from keras_hub.src.models.gemma.gemma_tokenizer import GemmaTokenizer
 from keras_hub.src.models.gpt2.gpt2_causal_lm_preprocessor import (
     GPT2CausalLMPreprocessor,
 )
@@ -43,3 +49,22 @@ def test_from_preset_errors(self):
         with self.assertRaises(ValueError):
             # No loading on a non-keras model.
             GPT2CausalLMPreprocessor.from_preset("hf://spacy/en_core_web_sm")
+
+    def test_export_supported_preprocessor(self):
+        proto = os.path.join(self.get_test_data_dir(), "gemma_export_vocab.spm")
+        tokenizer = GemmaTokenizer(proto=proto)
+        preprocessor = GemmaCausalLMPreprocessor(tokenizer=tokenizer)
+        export_path = os.path.join(self.get_temp_dir(), "export_preprocessor")
+        preprocessor.export_to_transformers(export_path)
+        # Basic check: tokenizer config exists
+        self.assertTrue(
+            os.path.exists(os.path.join(export_path, "tokenizer_config.json"))
+        )
+
+    def test_export_missing_tokenizer(self):
+        preprocessor = GemmaCausalLMPreprocessor(tokenizer=None)
+        export_path = os.path.join(
+            self.get_temp_dir(), "export_missing_tokenizer"
+        )
+        with self.assertRaises(ValueError):
+            preprocessor.export_to_transformers(export_path)

keras_hub/src/models/task_test.py

@@ -7,6 +7,12 @@
 
 from keras_hub.src.models.bert.bert_text_classifier import BertTextClassifier
 from keras_hub.src.models.causal_lm import CausalLM
+from keras_hub.src.models.gemma.gemma_backbone import GemmaBackbone
+from keras_hub.src.models.gemma.gemma_causal_lm import GemmaCausalLM
+from keras_hub.src.models.gemma.gemma_causal_lm_preprocessor import (
+    GemmaCausalLMPreprocessor,
+)
+from keras_hub.src.models.gemma.gemma_tokenizer import GemmaTokenizer
 from keras_hub.src.models.gpt2.gpt2_causal_lm import GPT2CausalLM
 from keras_hub.src.models.image_classifier import ImageClassifier
 from keras_hub.src.models.preprocessor import Preprocessor
@@ -171,3 +177,85 @@ def test_save_to_preset_custom_backbone_and_preprocessor(self):
         restored_task = ImageClassifier.from_preset(save_dir)
         actual = restored_task.predict(batch)
         self.assertAllClose(expected, actual)
+
+    def _create_gemma_for_export_tests(self):
+        proto = os.path.join(self.get_test_data_dir(), "gemma_export_vocab.spm")
+        tokenizer = GemmaTokenizer(proto=proto)
+        backbone = GemmaBackbone(
+            vocabulary_size=tokenizer.vocabulary_size(),
+            num_layers=2,
+            num_query_heads=4,
+            num_key_value_heads=1,
+            hidden_dim=512,
+            intermediate_dim=1024,
+            head_dim=128,
+        )
+        preprocessor = GemmaCausalLMPreprocessor(tokenizer=tokenizer)
+        causal_lm = GemmaCausalLM(backbone=backbone, preprocessor=preprocessor)
+        return causal_lm, preprocessor
+
+    def test_export_attached(self):
+        causal_lm, _ = self._create_gemma_for_export_tests()
+        export_path = os.path.join(self.get_temp_dir(), "export_attached")
+        causal_lm.export_to_transformers(export_path)
+        # Basic check: config and tokenizer files exist
+        self.assertTrue(
+            os.path.exists(os.path.join(export_path, "config.json"))
+        )
+        self.assertTrue(
+            os.path.exists(os.path.join(export_path, "tokenizer_config.json"))
+        )
+
+    def test_export_attached_with_lm_head(self):
+        # Since attached export always includes lm_head=True, this test verifies
+        # the same but explicitly notes it for coverage.
+        causal_lm, _ = self._create_gemma_for_export_tests()
+        export_path = os.path.join(
+            self.get_temp_dir(), "export_attached_lm_head"
+        )
+        causal_lm.export_to_transformers(export_path)
+        # Basic check: config and tokenizer files exist
+        self.assertTrue(
+            os.path.exists(os.path.join(export_path, "config.json"))
+        )
+        self.assertTrue(
+            os.path.exists(os.path.join(export_path, "tokenizer_config.json"))
+        )
+
+    def test_export_detached(self):
+        causal_lm, preprocessor = self._create_gemma_for_export_tests()
+        export_path_backbone = os.path.join(
+            self.get_temp_dir(), "export_detached_backbone"
+        )
+        export_path_preprocessor = os.path.join(
+            self.get_temp_dir(), "export_detached_preprocessor"
+        )
+        original_preprocessor = causal_lm.preprocessor
+        causal_lm.preprocessor = None
+        causal_lm.export_to_transformers(export_path_backbone)
+        causal_lm.preprocessor = original_preprocessor
+        preprocessor.export_to_transformers(export_path_preprocessor)
+        # Basic check: backbone has config, no tokenizer; preprocessor has
+        # tokenizer config
+        self.assertTrue(
+            os.path.exists(os.path.join(export_path_backbone, "config.json"))
+        )
+        self.assertFalse(
+            os.path.exists(
+                os.path.join(export_path_backbone, "tokenizer_config.json")
+            )
+        )
+        self.assertTrue(
+            os.path.exists(
+                os.path.join(export_path_preprocessor, "tokenizer_config.json")
+            )
+        )
+
+    def test_export_missing_tokenizer(self):
+        causal_lm, preprocessor = self._create_gemma_for_export_tests()
+        preprocessor.tokenizer = None
+        export_path = os.path.join(
+            self.get_temp_dir(), "export_missing_tokenizer"
+        )
+        with self.assertRaises(ValueError):
+            causal_lm.export_to_transformers(export_path)

keras_hub/src/tests/test_data/gemma_export_vocab.spm

@@ -261,3 +261,18 @@ class like `keras_hub.models.Tokenizer.from_preset()`, or from
         if cls.backbone_cls != backbone_cls:
             cls = find_subclass(preset, cls, backbone_cls)
         return loader.load_tokenizer(cls, config_file, **kwargs)
+
+    def export_to_transformers(self, path):
+        """Export the tokenizer to HuggingFace Transformers format.
+
+        This saves tokenizer assets in a format compatible with HuggingFace
+        Transformers.
+
+        Args:
+            path: str. Path to save the exported tokenizer.
+        """
+        from keras_hub.src.utils.transformers.export.hf_exporter import (
+            export_tokenizer,
+        )
+
+        export_tokenizer(self, path)

keras_hub/src/tests/test_data/llama_export_vocab.spm

@@ -6,6 +6,7 @@
 
 from keras_hub.src.models.albert.albert_tokenizer import AlbertTokenizer
 from keras_hub.src.models.bert.bert_tokenizer import BertTokenizer
+from keras_hub.src.models.gemma.gemma_tokenizer import GemmaTokenizer
 from keras_hub.src.models.gpt2.gpt2_tokenizer import GPT2Tokenizer
 from keras_hub.src.models.roberta.roberta_tokenizer import RobertaTokenizer
 from keras_hub.src.tests.test_case import TestCase
@@ -113,3 +114,24 @@ def test_save_to_preset(self, cls, preset_name, tokenizer_type):
         # Check config class.
         tokenizer_config = load_json(save_dir, TOKENIZER_CONFIG_FILE)
         self.assertEqual(cls, check_config_class(tokenizer_config))
+
+    def test_export_supported_tokenizer(self):
+        proto = os.path.join(self.get_test_data_dir(), "gemma_export_vocab.spm")
+        tokenizer = GemmaTokenizer(proto=proto)
+        export_path = os.path.join(self.get_temp_dir(), "export_tokenizer")
+        tokenizer.export_to_transformers(export_path)
+        # Basic check: tokenizer config exists
+        self.assertTrue(
+            os.path.exists(os.path.join(export_path, "tokenizer_config.json"))
+        )
+
+    def test_export_unsupported_tokenizer(self):
+        proto = os.path.join(self.get_test_data_dir(), "gemma_export_vocab.spm")
+
+        class UnsupportedTokenizer(GemmaTokenizer):
+            pass
+
+        tokenizer = UnsupportedTokenizer(proto=proto)
+        export_path = os.path.join(self.get_temp_dir(), "unsupported_tokenizer")
+        with self.assertRaises(ValueError):
+            tokenizer.export_to_transformers(export_path)