Model Export to liteRT

keras_hub/api/__init__.py

@@ -4,6 +4,7 @@
 since your modifications would be overwritten.
 """
 
+from keras_hub import export as export
 from keras_hub import layers as layers
 from keras_hub import metrics as metrics
 from keras_hub import models as models

keras_hub/api/export/__init__.py

@@ -0,0 +1,25 @@
+"""DO NOT EDIT.
+
+This file was autogenerated. Do not edit it by hand,
+since your modifications would be overwritten.
+"""
+
+from keras_hub.src.export.configs import (
+    CausalLMExporterConfig as CausalLMExporterConfig,
+)
+from keras_hub.src.export.configs import (
+    ImageClassifierExporterConfig as ImageClassifierExporterConfig,
+)
+from keras_hub.src.export.configs import (
+    ImageSegmenterExporterConfig as ImageSegmenterExporterConfig,
+)
+from keras_hub.src.export.configs import (
+    ObjectDetectorExporterConfig as ObjectDetectorExporterConfig,
+)
+from keras_hub.src.export.configs import (
+    Seq2SeqLMExporterConfig as Seq2SeqLMExporterConfig,
+)
+from keras_hub.src.export.configs import (
+    TextClassifierExporterConfig as TextClassifierExporterConfig,
+)
+from keras_hub.src.export.litert import LiteRTExporter as LiteRTExporter

keras_hub/src/export/__init__.py

@@ -0,0 +1,9 @@
+# Export base classes and configurations for advanced usage
+from keras_hub.src.export.base import KerasHubExporter
+from keras_hub.src.export.base import KerasHubExporterConfig
+from keras_hub.src.export.configs import CausalLMExporterConfig
+from keras_hub.src.export.configs import Seq2SeqLMExporterConfig
+from keras_hub.src.export.configs import TextClassifierExporterConfig
+from keras_hub.src.export.configs import get_exporter_config
+from keras_hub.src.export.litert import LiteRTExporter
+from keras_hub.src.export.litert import export_litert

keras_hub/src/export/base.py

@@ -0,0 +1,120 @@
+"""Base classes for Keras-Hub model exporters.
+
+This module provides the foundation for exporting Keras-Hub models to various
+formats. It defines the abstract base classes that all exporters must implement.
+"""
+
+from abc import ABC
+from abc import abstractmethod
+
+
+class KerasHubExporterConfig(ABC):
+    """Base configuration class for Keras-Hub model exporters.
+
+    This class defines the interface for exporter configurations that specify
+    how different types of Keras-Hub models should be exported.
+    """
+
+    # Model type this exporter handles (e.g., "causal_lm", "text_classifier")
+    MODEL_TYPE = None
+
+    # Expected input structure for this model type
+    EXPECTED_INPUTS = []
+
+    def __init__(self, model, **kwargs):
+        """Initialize the exporter configuration.
+
+        Args:
+            model: `keras.Model`. The Keras-Hub model to export.
+            **kwargs: Additional configuration parameters.
+        """
+        self.model = model
+        self.config_kwargs = kwargs
+        self._validate_model()
+
+    def _validate_model(self):
+        """Validate that the model is compatible with this exporter."""
+        if not self._is_model_compatible():
+            raise ValueError(
+                f"Model {self.model.__class__.__name__} is not compatible "
+                f"with {self.__class__.__name__}"
+            )
+
+    @abstractmethod
+    def _is_model_compatible(self):
+        """Check if the model is compatible with this exporter.
+
+        Returns:
+            `bool`. True if compatible, False otherwise
+        """
+        pass
+
+    @abstractmethod
+    def get_input_signature(self, sequence_length=None):
+        """Get the input signature for this model type.
+
+        Args:
+            sequence_length: `int` or `None`. Optional sequence length for
+                input tensors.
+
+        Returns:
+            `dict`. Dictionary mapping input names to tensor specifications.
+        """
+        pass
+
+
+class KerasHubExporter(ABC):
+    """Base class for Keras-Hub model exporters.
+
+    This class provides the common interface for exporting Keras-Hub models
+    to different formats (LiteRT, ONNX, etc.).
+    """
+
+    def __init__(self, config, **kwargs):
+        """Initialize the exporter.
+
+        Args:
+            config: `KerasHubExporterConfig`. Exporter configuration specifying
+                model type and parameters.
+            **kwargs: Additional exporter-specific parameters.
+        """
+        self.config = config
+        self.model = config.model
+        self.export_kwargs = kwargs
+
+    @abstractmethod
+    def export(self, filepath):
+        """Export the model to the specified filepath.
+
+        Args:
+            filepath: `str`. Path where to save the exported model.
+        """
+        pass
+
+    def _ensure_model_built(self, param=None):
+        """Ensure the model is properly built with correct input structure.
+
+        This method builds the model using model.build() with input shapes.
+        This creates the necessary variables and initializes the model structure
+        for export without needing dummy data.
+
+        Args:
+            param: `int` or `None`. Optional parameter for input signature
+                (e.g., sequence_length for text models, image_size for vision
+                models).
+        """
+        # Get input signature (returns dict of InputSpec objects)
+        input_signature = self.config.get_input_signature(param)
+
+        # Extract shapes from InputSpec objects
+        input_shapes = {}
+        for name, spec in input_signature.items():
+            if hasattr(spec, "shape"):
+                input_shapes[name] = spec.shape
+            else:
+                # Fallback for unexpected formats
+                input_shapes[name] = spec
+
+        # Build the model using shapes only (no actual data allocation)
+        # This creates variables and initializes the model structure
+        self.model.build(input_shape=input_shapes)

keras_hub/src/export/base_test.py

@@ -0,0 +1,197 @@
+"""Tests for base export classes."""
+
+import keras
+
+from keras_hub.src.export.base import KerasHubExporter
+from keras_hub.src.export.base import KerasHubExporterConfig
+from keras_hub.src.tests.test_case import TestCase
+
+
+class DummyExporterConfig(KerasHubExporterConfig):
+    """Dummy configuration for testing."""
+
+    MODEL_TYPE = "test_model"
+    EXPECTED_INPUTS = ["input_ids", "attention_mask"]
+    DEFAULT_SEQUENCE_LENGTH = 128
+
+    def __init__(self, model, compatible=True, **kwargs):
+        self.is_compatible = compatible
+        super().__init__(model, **kwargs)
+
+    def _is_model_compatible(self):
+        return self.is_compatible
+
+    def get_input_signature(self, sequence_length=None):
+        seq_len = sequence_length or self.DEFAULT_SEQUENCE_LENGTH
+        return {
+            "input_ids": keras.layers.InputSpec(
+                shape=(None, seq_len), dtype="int32"
+            ),
+            "attention_mask": keras.layers.InputSpec(
+                shape=(None, seq_len), dtype="int32"
+            ),
+        }
+
+
+class DummyExporter(KerasHubExporter):
+    """Dummy exporter for testing."""
+
+    def __init__(self, config, **kwargs):
+        super().__init__(config, **kwargs)
+        self.exported = False
+        self.export_path = None
+
+    def export(self, filepath):
+        self.exported = True
+        self.export_path = filepath
+        return filepath
+
+
+class KerasHubExporterConfigTest(TestCase):
+    """Tests for KerasHubExporterConfig base class."""
+
+    def test_init_with_compatible_model(self):
+        """Test initialization with a compatible model."""
+        model = keras.Sequential([keras.layers.Dense(10)])
+        config = DummyExporterConfig(model, compatible=True)
+
+        self.assertEqual(config.model, model)
+        self.assertEqual(config.MODEL_TYPE, "test_model")
+        self.assertEqual(
+            config.EXPECTED_INPUTS, ["input_ids", "attention_mask"]
+        )
+
+    def test_init_with_incompatible_model_raises_error(self):
+        """Test that incompatible model raises ValueError."""
+        model = keras.Sequential([keras.layers.Dense(10)])
+
+        with self.assertRaisesRegex(ValueError, "not compatible"):
+            DummyExporterConfig(model, compatible=False)
+
+    def test_get_input_signature_default_sequence_length(self):
+        """Test get_input_signature with default sequence length."""
+        model = keras.Sequential([keras.layers.Dense(10)])
+        config = DummyExporterConfig(model)
+
+        signature = config.get_input_signature()
+
+        self.assertIn("input_ids", signature)
+        self.assertIn("attention_mask", signature)
+        self.assertEqual(signature["input_ids"].shape, (None, 128))
+        self.assertEqual(signature["attention_mask"].shape, (None, 128))
+
+    def test_get_input_signature_custom_sequence_length(self):
+        """Test get_input_signature with custom sequence length."""
+        model = keras.Sequential([keras.layers.Dense(10)])
+        config = DummyExporterConfig(model)
+
+        signature = config.get_input_signature(sequence_length=256)
+
+        self.assertEqual(signature["input_ids"].shape, (None, 256))
+        self.assertEqual(signature["attention_mask"].shape, (None, 256))
+
+    def test_config_kwargs_stored(self):
+        """Test that additional kwargs are stored."""
+        model = keras.Sequential([keras.layers.Dense(10)])
+        config = DummyExporterConfig(
+            model, custom_param="value", another_param=42
+        )
+
+        self.assertEqual(config.config_kwargs["custom_param"], "value")
+        self.assertEqual(config.config_kwargs["another_param"], 42)
+
+
+class KerasHubExporterTest(TestCase):
+    """Tests for KerasHubExporter base class."""
+
+    def test_init_stores_config_and_model(self):
+        """Test that initialization stores config and model."""
+        model = keras.Sequential([keras.layers.Dense(10)])
+        config = DummyExporterConfig(model)
+        exporter = DummyExporter(config, verbose=True, custom_param="test")
+
+        self.assertEqual(exporter.config, config)
+        self.assertEqual(exporter.model, model)
+        self.assertEqual(exporter.export_kwargs["verbose"], True)
+        self.assertEqual(exporter.export_kwargs["custom_param"], "test")
+
+    def test_export_method_called(self):
+        """Test that export method can be called."""
+        model = keras.Sequential([keras.layers.Dense(10)])
+        config = DummyExporterConfig(model)
+        exporter = DummyExporter(config)
+
+        result = exporter.export("/tmp/test_model")
+
+        self.assertTrue(exporter.exported)
+        self.assertEqual(exporter.export_path, "/tmp/test_model")
+        self.assertEqual(result, "/tmp/test_model")
+
+    def test_ensure_model_built(self):
+        """Test _ensure_model_built method."""
+
+        class TestModel(keras.Model):
+            def __init__(self):
+                super().__init__()
+                self.dense = keras.layers.Dense(10)
+
+            def call(self, inputs):
+                return self.dense(inputs["input_ids"])
+
+        model = TestModel()
+        config = DummyExporterConfig(model)
+        exporter = DummyExporter(config)
+
+        # Model should not be built initially
+        self.assertFalse(model.built)
+
+        # Call _ensure_model_built
+        exporter._ensure_model_built()
+
+        # Model should now be built
+        self.assertTrue(model.built)
+
+    def test_ensure_model_built_with_custom_param(self):
+        """Test _ensure_model_built with custom sequence length."""
+
+        class TestModel(keras.Model):
+            def __init__(self):
+                super().__init__()
+                self.dense = keras.layers.Dense(10)
+
+            def call(self, inputs):
+                return self.dense(inputs["input_ids"])
+
+        model = TestModel()
+        config = DummyExporterConfig(model)
+        exporter = DummyExporter(config)
+
+        # Call with custom sequence length
+        exporter._ensure_model_built(param=512)
+
+        # Verify model is built
+        self.assertTrue(model.built)
+
+    def test_ensure_model_built_already_built_model(self):
+        """Test _ensure_model_built with already built model."""
+
+        class TestModel(keras.Model):
+            def __init__(self):
+                super().__init__()
+                self.dense = keras.layers.Dense(10)
+
+            def call(self, inputs):
+                return self.dense(inputs["input_ids"])
+
+        model = TestModel()
+        # Pre-build the model
+        model.build(input_shape={"input_ids": (None, 128)})
+
+        config = DummyExporterConfig(model)
+        exporter = DummyExporter(config)
+
+        # Should not raise an error for already built model
+        exporter._ensure_model_built()
+
+        # Model should still be built
+        self.assertTrue(model.built)