Adapt "Natural language image search with a Dual Encoder" Demo to Keras 3 and Keras Hub

[vinayak19th]

Summary

This PR updates the 'Natural language image search with a Dual Encoder' example to Keras 3 with Torch and Tensorflow Backend support for training, however, still using tf.data for dataloading and pre-processing.

Changelog

Major Changes:

1. Migrated all TensorFlow Hub functionality to KerasHub
2. Removed all tf.x functions to Keras 3 ops.x supporting full multi-backend functionality
3. Created multiple train_step implementations for the Tensorflow and Torch backends in the Dual Encoder with automatic selection of the correct implementation based on the backend.
4. Updated vision_encoder to use the better and more stable EfficientNetV2 models over the older Xception model.
5. Moved text preprocessing and tokenization out of the model and into the tf.Data pipeline for more efficient GPU VRAM usage and training speed.
6. Added additional blocks of code for increased clarity, such as a cell showing the processed dataset elements.
7. Fixed issues in dataset download and unzip file pathing due to changes in modern versions of the wget package.

Minor Changes:

1. Updated the markdown documentation to better match modern Keras 3 standards.
2. Converted tf.matmul-based inner product to the simplified and more efficient ops.einsum.
3. Added additional explanations to functions for increased clarity, such as an explanation of the different tf.data stages.

[gemini-code-assist]

Summary of Changes

Hello @vinayak19th, 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 significantly updates the 'Natural language image search with a Dual Encoder' demo to align with Keras 3, introducing multi-backend support for training with both TensorFlow and Torch. The changes focus on modernizing the underlying libraries, improving model architecture with a more advanced vision encoder, and optimizing the data processing pipeline for better performance and clarity. The overall goal is to provide a robust and up-to-date example for natural language image search within the Keras ecosystem.

Highlights

• Keras 3 Adaptation: The demo has been fully adapted to Keras 3, ensuring compatibility with both TensorFlow and Torch backends for training.
• Keras Hub Migration: All TensorFlow Hub functionalities have been migrated to KerasHub, streamlining dependencies and aligning with Keras 3 standards.
• Multi-Backend Training: The DualEncoder now includes separate train_step implementations for TensorFlow and Torch, with automatic selection based on the active Keras backend.
• Improved Vision Encoder: The vision encoder has been upgraded from Xception to the more stable and performant EfficientNetV2B3 model.
• Optimized Data Pipeline: Text preprocessing and tokenization have been moved from the model into the tf.data pipeline, enhancing GPU VRAM usage and training speed.
• Dataset Handling Fixes: Issues with dataset download and unzip file pathing, caused by changes in modern wget package versions, have been resolved.
• Code Clarity and Efficiency: Added code blocks for clarity (e.g., showing processed dataset elements) and converted tf.matmul operations to ops.einsum for improved efficiency.

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.

Footnotes

1. 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 successfully updates the "Natural language image search with a Dual Encoder" demo to Keras 3, including migrating from TensorFlow Hub to Keras Hub and adding multi-backend support. The changes are well-structured, with notable improvements like upgrading the vision encoder to EfficientNetV2 and moving text preprocessing into the tf.data pipeline for better performance. My feedback includes a critical suggestion to reintroduce shuffling in the data pipeline to ensure proper model training, along with several medium-severity recommendations to improve code clarity and maintainability by removing commented-out code, unused variables, and replacing hardcoded values with defined constants.

[vinayak19th]

/gemini review

[gemini-code-assist]

Code Review

This is a great pull request that successfully updates the example to Keras 3 and Keras Hub, enabling multi-backend support and using modern components. The code is much clearer and follows current best practices. I've found a critical syntax error in a print statement that needs to be fixed. I also have some high-severity feedback regarding the dataset download and extraction logic, specifically about cleaning up large temporary files and ensuring path consistency. Additionally, I've included a few medium-severity suggestions to improve robustness and add support for the JAX backend to make it fully Keras 3 compatible. Overall, excellent work on this migration!

[gemini-code-assist]

This logic for selecting train_step is great for multi-backend support. However, it only covers TensorFlow and Torch. To make this example fully compatible with Keras 3, you could also add support for the JAX backend. This would require implementing a train_step_jax method.

[vinayak19th]

Not sure how to implement in Jax fully yet