KerasHub

KerasHub is designed around several core principles to ensure it serves both beginners and experts in the NLP field:

• High-level abstraction: Even simple tasks should be packaged as one-liners to maximize ease of use.
• Balance of ease and flexibility: Simple workflows should be easy to implement, but the library must provide a "go deeper" path for advanced users to customize components.
• Modular building blocks: The scope includes all necessary NLP components, such as data loading, augmentation, model building, evaluation metrics, and visualization utilities.
• Language Agnostic Design: Workflows should prioritize multi-lingual support and avoid language-specific logic (like stemming) that requires per-language rewrites.

KerasHub aims to be self-contained, relying primarily on Keras, NumPy, TensorFlow, and Tensorflow Text. It avoids heavy external NLP dependencies like NLTK or spaCy for preprocessing.

In cases where an external dependency is required for a specific metric or tokenizer (e.g., rouge_score), KerasHub uses an optional import pattern. If the dependency is missing, the library will raise an ImportError with specific installation instructions rather than failing at the module level.

try:
    import rouge_score
except ImportError:
    rouge_score = None

class Rouge(keras.metrics.Metric):
    def __init__(self):
        if rouge_score is None:
            raise ImportError(
                "ROUGE metric requires the `rouge_score` package."
                "Please install it with `pip install rouge_score`."
            )

To ensure high performance and efficiency, KerasHub components (layers, metrics, and tokenizers) follow these computational guidelines:

• TensorFlow Graph Compatibility: All components should be compatible with @tf.function. They should rely on tf.strings and tf.text operations to ensure string manipulation happens within the TensorFlow graph.
• XLA Compilation: Trainable modeling components should be designed to be XLA compilable (e.g., compatible with tf.function(jit_compile=True)) to leverage significant performance gains.
• tf.data Integration: Preprocessing tools must be runnable inside a tf.data pipeline. Augmentations should be dynamic (performed on-the-fly during training) rather than precomputed. Preprocessing layers should support both batched and unbatched data inputs.

When naming layers and models, follow these rules:

• Acronyms: Capitalize all acronyms (e.g., LSTM, KLDivergence, GPT2, XLMRoberta).
• Pronounceable Acronyms: If an abbreviation is common and pronounceable, treat it as a standalone word (e.g., Bert, Deberta).
• File Names: Use snake_case. Treat acronyms as a single segment. For example, XLMRoberta should be in a file named xlm_roberta.py.
• Public Classes: Ideally, keep publicly documented classes in their own files, matching the class name to the filename (e.g., BertClassifier in bert_classifier.py).

To install the latest stable release of KerasHub with Keras 3, use pip. Note that installing KerasHub will automatically pull in TensorFlow to support the tf.data API for preprocessing, though training can still be performed on other backends like JAX or PyTorch.

To install the latest nightly builds for both KerasHub and Keras, use the nightly package.

# Install latest stable release
pip install --upgrade keras-hub

# Install latest nightly changes
pip install --upgrade keras-hub-nightly

This optional stage extends the model to specific downstream tasks (e.g., Masked LM, Classification). The following files and verification steps are required:

• xx/xx_<task>.py: The task model implementation (e.g., a classifier head).
• xx/xx_<task>_preprocessor.py: The preprocessor that prepares inputs suitable for the task model.
• xx/xx_<task>_test.py and xx/xx_<task>_preprocessor_test.py: Unit tests for both the task model and the preprocessor.
• Verification: A Colab notebook link must be included in the PR description. This notebook must demonstrate that the output of your preprocessor matches the output of the original preprocessor.

The second step in model contribution is adding the tokenizer. Most text models use subword tokenizers (WordPiece, SentencePiece, BPE).

Implementation Steps:

1. Inheritance: Inherit from an existing KerasHub base tokenizer class (e.g., keras_hub.tokenizers.WordPieceTokenizer).
2. Special Tokens: You must add special tokens (e.g., beginning-of-sequence, end-of-sequence, mask, pad) as member attributes to the tokenizer class. These attributes are used by preprocessor layers.
3. Unit Tests: Create tests using a dummy vocabulary to verify tokenization and detokenization processes.

# Example pattern
class MyModelTokenizer(keras_hub.tokenizers.WordPieceTokenizer):
    def __init__(self, ...):
        super().__init__(...)
        self.cls_token = '...' # Special tokens as member attributes
        self.sep_token = '...'

To contribute the first PR for a new model, you must implement the backbone architecture. The following files and verification steps are required:

• xx/xx_backbone.py: Contains the model graph implementation.
• xx/xx_backbone_test.py: Contains unit tests for the backbone.
• Verification: A Colab notebook link must be included in the PR description. This notebook must demonstrate that the outputs of your implemented backbone match the outputs of the original source model.