$ show keras-team/keras-hub
KerasHub is a production-grade Keras 3 library providing pretrained model architectures, layers, tokenizers, metrics, and sampling utilities for NLP, vision, multimodal, and audio tasks. The codebase is organized around a consistent Backbone → Task → Preprocessor pattern, where each model family (e.g., BERT, Gemma, ViT, Whisper) defines a backbone, a preprocessor (tokenizer + optional media converter), and task-specific heads (CausalLM, ImageClassifier, etc.). All models expose a unified `from_preset()` API for loading pretrained weights hosted on Kaggle Models, and `save_to_preset()` for export. The library includes extensive checkpoint conversion tooling from HuggingFace and TIMM, export to HuggingFace format and OpenVINO, and integration with vLLM for serving. While predominantly an application library, it has strong data/research character through its pretrained weight infrastructure, benchmark scripts, and research-oriented architectures (e.g., FNet, RWKV7, BLIP-2).
The LLM identified 5 files worth highlighting for reuse:
Implements a `ReversibleEmbedding` layer that shares weights between the token embedding and output projection (tied embeddings). This is a self-contained pattern used across many backbones (BERT, GPT, Llama, etc.) to reduce parameter count. The layer supports mixed precision and is fully configurable via `get_config()`. It demonstrates non-obvious skill in handling gradient scaling and embedding reuse within the Keras functional API.
A Keras `PreprocessingLayer` that packs multiple text segments into a single token sequence with optional start/end tokens, segment IDs, and padding. It is used by most text classifier and masked LM preprocessors. The layer is fully parameterized (max length, truncation strategy, special token IDs) and can be used standalone or within `tf.data` pipelines. Its design isolates the packing logic from tokenization, making it reusable across model families.
Provides a base class for ROUGE metrics (ROUGE-N, ROUGE-L) as Keras metrics. It manages state for true/false positives and handles ragged string tensors from TF `tf.data` pipelines. The implementation correctly handles precision/recall/F1 computation and serialization. It demonstrates a clean pattern for building custom Keras metrics that work across backends, though it depends on the external `rouge_score` package for tokenization.
Implements a generic beam search sampler for autoregressive generation. It is decoupled from any specific model architecture and only requires a `next_token_fn`. Supports variable-length beams, early stopping, and length penalty. The implementation uses `keras.ops` for backend-agnostic tensor operations and includes caching support for efficient decoding. This is a self-contained algorithm that can be extracted and reused in any text generation pipeline.
Provides `PresetLoader` and `PresetSaver` classes that handle downloading, caching, and loading pretrained model weights from Kaggle Models or local directories. It supports multi-file checkpoints, sharded weights, and safetensors. The loader abstracts away the storage backend and provides a uniform interface for `from_preset()` across all models. This utility is critical for the library's weight distribution system and can be adapted for other Keras-based model hubs.