Model Export to liteRT #2405

pctablet505 · 2025-09-17T09:25:59Z

Added support for model export, currently able to convert gemma3, llama3.2, gpt2 models, and verified numerics also.

Colab Notebook

(https://colab.research.google.com/gist/pctablet505/45a48c42fa91cc27995cdaefda57cb28/model-export.ipynb)

This reverts commit 62d2484.

This reverts commit de830b1.

export working 1st commit

Refactored exporter and registry logic for better type safety and error handling. Improved input signature methods in config classes by extracting sequence length logic. Enhanced LiteRT exporter with clearer verbose handling and stricter error reporting. Registry now conditionally registers LiteRT exporter and extends export method only if dependencies are available.

gemini-code-assist

Summary of Changes

Hello @pctablet505, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces a comprehensive and extensible framework for exporting Keras-Hub models to various formats, with an initial focus on LiteRT. The system is designed to seamlessly integrate with Keras-Hub's model architecture, particularly by addressing the unique challenge of handling dictionary-based model inputs during the export process. This enhancement significantly improves the deployability of Keras-Hub models by providing a standardized and robust export pipeline, alongside crucial compatibility fixes for TensorFlow's SavedModel/TFLite export mechanisms.

Highlights

New Model Export Framework: Introduced a new, extensible framework for exporting Keras-Hub models, designed to support various formats and model types.
LiteRT Export Support: Added specific support for exporting Keras-Hub models to the LiteRT format, verified for models like gemma3, llama3.2, and gpt2.
Registry-Based Configuration: Implemented an ExporterRegistry to manage and retrieve appropriate exporter configurations and exporters based on model type and target format.
Input Handling for Keras-Hub Models: Developed a KerasHubModelWrapper to seamlessly convert Keras-Hub's dictionary-based inputs to the list-based inputs expected by the underlying Keras LiteRT exporter.
TensorFlow Export Compatibility: Added compatibility shims (_get_save_spec and _trackable_children) to Keras-Hub Backbone models to ensure proper functioning with TensorFlow's SavedModel and TFLite export utilities.
Automated Export Method Extension: The Task class in Keras-Hub models is now automatically extended with an export method, simplifying the model export process for users.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	`/gemini review`	Performs a code review for the current pull request in its current state.
Pull Request Summary	`/gemini summary`	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	`/gemini help`	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

gemini-code-assist

Code Review

This pull request introduces a significant new feature: model exporting to liteRT. The implementation is well-structured, using a modular and extensible registry pattern. However, there are several areas that require attention. The most critical issue is the complete absence of tests for the new export functionality, which is a direct violation of the repository's style guide stating that testing is non-negotiable. Additionally, I've identified a critical bug in the error handling logic within the lite_rt.py exporter that includes unreachable code. There are also several violations of the style guide regarding the use of type hints in function signatures across all new files. I've provided specific comments and suggestions to address these points, which should help improve the robustness, maintainability, and compliance of this new feature.

gemini-code-assist · 2025-09-17T09:28:20Z

keras_hub/src/export/lite_rt.py

+        try:
+            # Export using the Keras exporter
+            keras_exporter.export(filepath)
+
+            if self.verbose:
+                print(f"Export completed successfully to: {filepath}.tflite")
+
+        except Exception as e:
+            raise RuntimeError(f"LiteRT export failed: {e}") from e
+            keras_exporter.export(filepath)
+
+            if self.verbose:
+                print(f"✅ Export completed successfully!")
+                print(f"📁 Model saved to: {filepath}.tflite")
+
+        except Exception as e:
+            if self.verbose:
+                print(f"❌ Export failed: {e}")
+            raise


The try...except block for the export process contains a bug. The code from line 102 is unreachable due to the raise statement on line 101. Additionally, having two consecutive except Exception as e: blocks is a syntax error. The error handling logic should be corrected to properly handle exceptions and report success or failure.

Suggested change

try:

# Export using the Keras exporter

keras_exporter.export(filepath)

if self.verbose:

print(f"Export completed successfully to: {filepath}.tflite")

except Exception as e:

raise RuntimeError(f"LiteRT export failed: {e}") from e

keras_exporter.export(filepath)

if self.verbose:

print(f"✅ Export completed successfully!")

print(f"📁 Model saved to: {filepath}.tflite")

except Exception as e:

if self.verbose:

print(f"❌ Export failed: {e}")

raise

try:

# Export using the Keras exporter

keras_exporter.export(filepath)

if self.verbose:

print(f"Export completed successfully to: {filepath}.tflite")

except Exception as e:

if self.verbose:

print(f"❌ Export failed: {e}")

raise RuntimeError(f"LiteRT export failed: {e}") from e

gemini-code-assist · 2025-09-17T09:28:21Z

keras_hub/src/export/base.py

+            )
+
+    @abstractmethod
+    def _is_model_compatible(self) -> bool:


The style guide specifies that type hints should not be used in function signatures ¹. Instead, type information should be provided in the Args section of the docstring. This rule is violated throughout the new files. For example, _is_model_compatible(self) -> bool:. Please remove type hints from all function signatures in this file and ensure the types are documented in the docstrings.

Style Guide References

Footnotes

KerasHub does not use type hints in function signatures. Type information should be provided in the docstring. ↩

gemini-code-assist · 2025-09-17T09:28:21Z

keras_hub/src/export/configs.py

+    EXPECTED_INPUTS = ["token_ids", "padding_mask"]
+    DEFAULT_SEQUENCE_LENGTH = 128
+
+    def _is_model_compatible(self) -> bool:


Similar to other new files, this file uses type hints in function signatures (e.g., _is_model_compatible(self) -> bool:), which is against the style guide ¹. Please remove the type hints from function signatures and move the type information to the docstrings.

Style Guide References

Footnotes

KerasHub does not use type hints in function signatures. Type information should be provided in the docstring. ↩

gemini-code-assist · 2025-09-17T09:28:21Z

keras_hub/src/export/configs.py

+    def _get_sequence_length(self) -> int:
+        """Get sequence length from model or use default."""
+        if hasattr(self.model, 'preprocessor') and self.model.preprocessor:
+            return getattr(self.model.preprocessor, 'sequence_length', self.DEFAULT_SEQUENCE_LENGTH)
+        return self.DEFAULT_SEQUENCE_LENGTH


The _get_sequence_length method is duplicated across CausalLMExporterConfig, TextClassifierExporterConfig, Seq2SeqLMExporterConfig, and TextModelExporterConfig. To improve maintainability and reduce code duplication, this method should be moved to the base class KerasHubExporterConfig in keras_hub/src/export/base.py.

gemini-code-assist · 2025-09-17T09:28:21Z

keras_hub/src/export/lite_rt.py

+    def __init__(self, config: KerasHubExporterConfig, 
+                 max_sequence_length: Optional[int] = None,
+                 aot_compile_targets: Optional[list] = None,
+                 verbose: bool = False,
+                 **kwargs):


This file uses type hints in function signatures (e.g., __init__(self, config: KerasHubExporterConfig, ...)), which is against the style guide ¹. Please remove the type hints and move the type information to the docstrings.

Style Guide References

Footnotes

KerasHub does not use type hints in function signatures. Type information should be provided in the docstring. ↩

gemini-code-assist · 2025-09-17T09:28:21Z

keras_hub/src/export/registry.py

+        pass
+
+
+def export_model(model, filepath: str, format: str = "lite_rt", **kwargs):


This file uses type hints in function signatures (e.g., export_model(model, filepath: str, ...)), which is against the style guide ¹. Please remove the type hints and move the type information to the docstrings.

Style Guide References

Footnotes

KerasHub does not use type hints in function signatures. Type information should be provided in the docstring. ↩

gemini-code-assist · 2025-09-17T09:28:21Z

keras_hub/src/export/registry.py

+        **kwargs: Additional arguments passed to the exporter
+    """
+    # Ensure registry is initialized
+    initialize_export_registry()


The initialize_export_registry() function is called at the module level on line 125, ensuring it runs once upon import. Calling it again within the export_model function is redundant and can be removed for minor performance improvement and cleaner code.

gemini-code-assist · 2025-09-17T09:28:21Z

keras_hub/src/models/__init__.py

+except ImportError as e:
+    print(f"⚠️  Failed to import Keras-Hub export functionality: {e}")


Using print for warnings is not ideal as it doesn't give developers control over the message. It's better practice to use warnings.warn to log this failure. This allows users of the library to filter or redirect warnings as needed.

Suggested change

except ImportError as e:

print(f"⚠️ Failed to import Keras-Hub export functionality: {e}")

except ImportError as e:

import warnings

warnings.warn(f"⚠️ Failed to import Keras-Hub export functionality: {e}")

Introduces the keras_hub.api.export submodule and updates the main API to expose it. The new export module imports various exporter configs and functions from the internal export package, making them available through the public API.

Added ImageClassifierExporterConfig, ImageSegmenterExporterConfig, and ObjectDetectorExporterConfig to the export API. Improved input shape inference and dummy input generation for image-related exporter configs. Refactored LiteRTExporter to better handle model type checks and input signature logic, with improved error handling for input mapping.

Moved the 'import keras' statement to the top of the module and removed redundant local imports within class methods. This improves code clarity and avoids repeated imports.

Deleted the debug_object_detection.py script, which was used for testing object detection model outputs and export issues. This cleanup removes unused debugging code from the repository.

Renames all references of 'LiteRT' to 'Litert' across the codebase, including file names, class names, and function names. Updates exporter registry and API imports to use the new 'litert' naming. Also improves image model exporter configs to dynamically determine input dtype from the model, enhancing flexibility for different input types. Adds support for ImageSegmenter model type detection in the exporter registry.

pctablet505 and others added 11 commits September 1, 2025 19:11

Update backbone.py

087b9b2

Update backbone.py

de830b1

Update task.py

62d2484

Revert "Update task.py"

3b71125

This reverts commit 62d2484.

Revert "Update backbone.py"

3d453ff

This reverts commit de830b1.

export

92b1254

export working 1st commit

refactoring

e46241d

refactor

6e970e2

Update registry.py

15ad9f3

Merge branch 'keras-team:master' into export

901c233

github-actions bot added the Gemma Gemma model specific issues label Sep 17, 2025

gemini-code-assist bot reviewed Sep 17, 2025

View reviewed changes

pctablet505 added 8 commits September 22, 2025 11:06

reformat

442fdd3

Add export submodule to keras_hub API

5446e2a

Introduces the keras_hub.api.export submodule and updates the main API to expose it. The new export module imports various exporter configs and functions from the internal export package, making them available through the public API.

reformat

5c31d88

now supporting export for objectDetectors

3290d42

Refactor: move keras import to module level

8df5a75

Moved the 'import keras' statement to the top of the module and removed redundant local imports within class methods. This improves code clarity and avoids repeated imports.

Remove debug_object_detection.py script

759d223

Deleted the debug_object_detection.py script, which was used for testing object detection model outputs and export issues. This cleanup removes unused debugging code from the repository.

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Model Export to liteRT #2405

Model Export to liteRT #2405

pctablet505 commented Sep 17, 2025

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

gemini-code-assist bot Sep 17, 2025

Uh oh!

gemini-code-assist bot Sep 17, 2025

Uh oh!

gemini-code-assist bot Sep 17, 2025

Uh oh!

gemini-code-assist bot Sep 17, 2025

Uh oh!

gemini-code-assist bot Sep 17, 2025

Uh oh!

gemini-code-assist bot Sep 17, 2025

Uh oh!

gemini-code-assist bot Sep 17, 2025

Uh oh!

gemini-code-assist bot Sep 17, 2025

Uh oh!

Uh oh!

		pass


		def export_model(model, filepath: str, format: str = "lite_rt", **kwargs):

		except ImportError as e:
		print(f"⚠️ Failed to import Keras-Hub export functionality: {e}")

Model Export to liteRT #2405

Are you sure you want to change the base?

Model Export to liteRT #2405

Conversation

pctablet505 commented Sep 17, 2025

Colab Notebook

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist bot Sep 17, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Sep 17, 2025

Choose a reason for hiding this comment

Style Guide References

Footnotes

Uh oh!

gemini-code-assist bot Sep 17, 2025

Choose a reason for hiding this comment

Style Guide References

Footnotes

Uh oh!

gemini-code-assist bot Sep 17, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Sep 17, 2025

Choose a reason for hiding this comment

Style Guide References

Footnotes

Uh oh!

gemini-code-assist bot Sep 17, 2025

Choose a reason for hiding this comment

Style Guide References

Footnotes

Uh oh!

gemini-code-assist bot Sep 17, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Sep 17, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!