keras-hub-nightly 0.15.0.dev20240823171555__py3-none-any.whl → 0.15.0.dev20240911134614__py3-none-any.whl

keras_hub/src/models/f_net/f_net_masked_lm_preprocessor.py

@@ -13,17 +13,16 @@
 # limitations under the License.
 
 import keras
-from absl import logging
 
 from keras_hub.src.api_export import keras_hub_export
-from keras_hub.src.layers.preprocessing.masked_lm_mask_generator import (
-    MaskedLMMaskGenerator,
-)
-from keras_hub.src.models.f_net.f_net_preprocessor import FNetPreprocessor
+from keras_hub.src.models.f_net.f_net_backbone import FNetBackbone
+from keras_hub.src.models.f_net.f_net_tokenizer import FNetTokenizer
+from keras_hub.src.models.masked_lm_preprocessor import MaskedLMPreprocessor
+from keras_hub.src.utils.tensor_utils import preprocessing_function
 
 
 @keras_hub_export("keras_hub.models.FNetMaskedLMPreprocessor")
-class FNetMaskedLMPreprocessor(FNetPreprocessor):
+class FNetMaskedLMPreprocessor(MaskedLMPreprocessor):
     """FNet preprocessing for the masked language modeling task.
 
     This preprocessing layer will prepare inputs for a masked language modeling
@@ -119,78 +118,13 @@ class FNetMaskedLMPreprocessor(FNetPreprocessor):
     ```
     """
 
-    def __init__(
-        self,
-        tokenizer,
-        sequence_length=512,
-        truncate="round_robin",
-        mask_selection_rate=0.15,
-        mask_selection_length=96,
-        mask_token_rate=0.8,
-        random_token_rate=0.1,
-        **kwargs,
-    ):
-        super().__init__(
-            tokenizer,
-            sequence_length=sequence_length,
-            truncate=truncate,
-            **kwargs,
-        )
-        self.mask_selection_rate = mask_selection_rate
-        self.mask_selection_length = mask_selection_length
-        self.mask_token_rate = mask_token_rate
-        self.random_token_rate = random_token_rate
-        self.masker = None
-
-    def build(self, input_shape):
-        super().build(input_shape)
-        # Defer masker creation to `build()` so that we can be sure tokenizer
-        # assets have loaded when restoring a saved model.
-        self.masker = MaskedLMMaskGenerator(
-            mask_selection_rate=self.mask_selection_rate,
-            mask_selection_length=self.mask_selection_length,
-            mask_token_rate=self.mask_token_rate,
-            random_token_rate=self.random_token_rate,
-            vocabulary_size=self.tokenizer.vocabulary_size(),
-            mask_token_id=self.tokenizer.mask_token_id,
-            unselectable_token_ids=[
-                self.tokenizer.cls_token_id,
-                self.tokenizer.sep_token_id,
-                self.tokenizer.pad_token_id,
-            ],
-        )
-
-    def get_config(self):
-        config = super().get_config()
-        config.update(
-            {
-                "mask_selection_rate": self.mask_selection_rate,
-                "mask_selection_length": self.mask_selection_length,
-                "mask_token_rate": self.mask_token_rate,
-                "random_token_rate": self.random_token_rate,
-            }
-        )
-        return config
+    backbone_cls = FNetBackbone
+    tokenizer_cls = FNetTokenizer
 
+    @preprocessing_function
     def call(self, x, y=None, sample_weight=None):
-        if y is not None or sample_weight is not None:
-            logging.warning(
-                f"{self.__class__.__name__} generates `y` and `sample_weight` "
-                "based on your input data, but your data already contains `y` "
-                "or `sample_weight`. Your `y` and `sample_weight` will be "
-                "ignored."
-            )
-        x = super().call(x)
-        token_ids, segment_ids = (
-            x["token_ids"],
-            x["segment_ids"],
-        )
-        masker_outputs = self.masker(token_ids)
-        x = {
-            "token_ids": masker_outputs["token_ids"],
-            "segment_ids": segment_ids,
-            "mask_positions": masker_outputs["mask_positions"],
-        }
-        y = masker_outputs["mask_ids"]
-        sample_weight = masker_outputs["mask_weights"]
+        output = super().call(x, y=y, sample_weight=sample_weight)
+        x, y, sample_weight = keras.utils.unpack_x_y_sample_weight(output)
+        # FNet has not padding mask.
+        del x["padding_mask"]
         return keras.utils.pack_x_y_sample_weight(x, y, sample_weight)

keras_hub/src/models/f_net/{f_net_classifier.py → f_net_text_classifier.py}

@@ -16,14 +16,16 @@
 import keras
 
 from keras_hub.src.api_export import keras_hub_export
-from keras_hub.src.models.classifier import Classifier
 from keras_hub.src.models.f_net.f_net_backbone import FNetBackbone
 from keras_hub.src.models.f_net.f_net_backbone import f_net_kernel_initializer
-from keras_hub.src.models.f_net.f_net_preprocessor import FNetPreprocessor
+from keras_hub.src.models.f_net.f_net_text_classifier_preprocessor import (
+    FNetTextClassifierPreprocessor,
+)
+from keras_hub.src.models.text_classifier import TextClassifier
 
 
-@keras_hub_export("keras_hub.models.FNetClassifier")
-class FNetClassifier(Classifier):
+@keras_hub_export("keras_hub.models.FNetTextClassifier")
+class FNetTextClassifier(TextClassifier):
     """An end-to-end f_net model for classification tasks.
 
     This model attaches a classification head to a
@@ -42,7 +44,7 @@ class FNetClassifier(Classifier):
     Args:
         backbone: A `keras_hub.models.FNetBackbone` instance.
         num_classes: int. Number of classes to predict.
-        preprocessor: A `keras_hub.models.FNetPreprocessor` or `None`. If
+        preprocessor: A `keras_hub.models.FNetTextClassifierPreprocessor` or `None`. If
             `None`, this model will not apply preprocessing, and inputs should
             be preprocessed before calling the model.
         activation: Optional `str` or callable. The
@@ -61,7 +63,7 @@ class FNetClassifier(Classifier):
     labels = [0, 3]
 
     # Pretrained classifier.
-    classifier = keras_hub.models.FNetClassifier.from_preset(
+    classifier = keras_hub.models.FNetTextClassifier.from_preset(
         "f_net_base_en",
         num_classes=4,
     )
@@ -89,7 +91,7 @@ class FNetClassifier(Classifier):
     labels = [0, 3]
 
     # Pretrained classifier without preprocessing.
-    classifier = keras_hub.models.FNetClassifier.from_preset(
+    classifier = keras_hub.models.FNetTextClassifier.from_preset(
         "f_net_base_en",
         num_classes=4,
         preprocessor=None,
@@ -99,7 +101,7 @@ class FNetClassifier(Classifier):
     """
 
     backbone_cls = FNetBackbone
-    preprocessor_cls = FNetPreprocessor
+    preprocessor_cls = FNetTextClassifierPreprocessor
 
     def __init__(
         self,

keras_hub/src/models/f_net/{f_net_preprocessor.py → f_net_text_classifier_preprocessor.py}

@@ -16,18 +16,21 @@
 import keras
 
 from keras_hub.src.api_export import keras_hub_export
-from keras_hub.src.layers.preprocessing.multi_segment_packer import (
-    MultiSegmentPacker,
-)
+from keras_hub.src.models.f_net.f_net_backbone import FNetBackbone
 from keras_hub.src.models.f_net.f_net_tokenizer import FNetTokenizer
-from keras_hub.src.models.preprocessor import Preprocessor
-from keras_hub.src.utils.keras_utils import (
-    convert_inputs_to_list_of_tensor_segments,
+from keras_hub.src.models.text_classifier_preprocessor import (
+    TextClassifierPreprocessor,
 )
+from keras_hub.src.utils.tensor_utils import preprocessing_function
 
 
-@keras_hub_export("keras_hub.models.FNetPreprocessor")
-class FNetPreprocessor(Preprocessor):
+@keras_hub_export(
+    [
+        "keras_hub.models.FNetTextClassifierPreprocessor",
+        "keras_hub.models.FNetPreprocessor",
+    ]
+)
+class FNetTextClassifierPreprocessor(TextClassifierPreprocessor):
     """An FNet preprocessing layer which tokenizes and packs inputs.
 
     This preprocessing layer will do three things:
@@ -68,7 +71,7 @@ class FNetPreprocessor(Preprocessor):
 
     Directly calling the from_preset().
     ```python
-    preprocessor = keras_hub.models.FNetPreprocessor.from_preset(
+    preprocessor = keras_hub.models.TextClassifierPreprocessor.from_preset(
         "f_net_base_en"
     )
 
@@ -87,7 +90,7 @@ class FNetPreprocessor(Preprocessor):
 
     Mapping with `tf.data.Dataset`.
     ```python
-    preprocessor = keras_hub.models.FNetPreprocessor.from_preset(
+    preprocessor = keras_hub.models.TextClassifierPreprocessor.from_preset(
         "f_net_base_en"
     )
     first = tf.constant(["The quick brown fox jumped.", "Call me Ishmael."])
@@ -118,60 +121,13 @@ class FNetPreprocessor(Preprocessor):
     ```
     """
 
+    backbone_cls = FNetBackbone
     tokenizer_cls = FNetTokenizer
 
-    def __init__(
-        self,
-        tokenizer,
-        sequence_length=512,
-        truncate="round_robin",
-        **kwargs,
-    ):
-        super().__init__(**kwargs)
-        self.tokenizer = tokenizer
-        self.packer = None
-        self.truncate = truncate
-        self.sequence_length = sequence_length
-
-    def build(self, input_shape):
-        # Defer packer creation to `build()` so that we can be sure tokenizer
-        # assets have loaded when restoring a saved model.
-        self.packer = MultiSegmentPacker(
-            start_value=self.tokenizer.cls_token_id,
-            end_value=self.tokenizer.sep_token_id,
-            pad_value=self.tokenizer.pad_token_id,
-            truncate=self.truncate,
-            sequence_length=self.sequence_length,
-        )
-        self.built = True
-
-    def get_config(self):
-        config = super().get_config()
-        config.update(
-            {
-                "sequence_length": self.sequence_length,
-                "truncate": self.truncate,
-            }
-        )
-        return config
-
+    @preprocessing_function
     def call(self, x, y=None, sample_weight=None):
-        x = convert_inputs_to_list_of_tensor_segments(x)
-        x = [self.tokenizer(segment) for segment in x]
-        token_ids, segment_ids = self.packer(x)
-        x = {
-            "token_ids": token_ids,
-            "segment_ids": segment_ids,
-        }
+        # FNet has not padding mask.
+        output = super().call(x, y=y, sample_weight=sample_weight)
+        x, y, sample_weight = keras.utils.unpack_x_y_sample_weight(output)
+        del x["padding_mask"]
         return keras.utils.pack_x_y_sample_weight(x, y, sample_weight)
-
-    @property
-    def sequence_length(self):
-        """The padded length of model input sequences."""
-        return self._sequence_length
-
-    @sequence_length.setter
-    def sequence_length(self, value):
-        self._sequence_length = value
-        if self.packer is not None:
-            self.packer.sequence_length = value

keras_hub/src/models/f_net/f_net_tokenizer.py

@@ -14,12 +14,18 @@
 
 
 from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.models.f_net.f_net_backbone import FNetBackbone
 from keras_hub.src.tokenizers.sentence_piece_tokenizer import (
     SentencePieceTokenizer,
 )
 
 
-@keras_hub_export("keras_hub.models.FNetTokenizer")
+@keras_hub_export(
+    [
+        "keras_hub.tokenizers.FNetTokenizer",
+        "keras_hub.models.FNetTokenizer",
+    ]
+)
 class FNetTokenizer(SentencePieceTokenizer):
     """FNet tokenizer layer based on SentencePiece.
 
@@ -29,10 +35,6 @@ class FNetTokenizer(SentencePieceTokenizer):
     FNet models and provides a `from_preset()` method to automatically
     download a matching vocabulary for a FNet preset.
 
-    This tokenizer does not provide truncation or padding of inputs. It can be
-    combined with a `keras_hub.models.FNetPreprocessor` layer for input
-    packing.
-
     If input is a batch of strings (rank > 0), the layer will output a
     `tf.RaggedTensor` where the last dimension of the output is ragged.
 
@@ -61,35 +63,15 @@ class FNetTokenizer(SentencePieceTokenizer):
     ```
     """
 
+    backbone_cls = FNetBackbone
+
     def __init__(self, proto, **kwargs):
-        self.cls_token = "[CLS]"
-        self.sep_token = "[SEP]"
-        self.pad_token = "<pad>"
-        self.mask_token = "[MASK]"
+        self._add_special_token("[CLS]", "cls_token")
+        self._add_special_token("[SEP]", "sep_token")
+        self._add_special_token("<pad>", "pad_token")
+        self._add_special_token("[MASK]", "mask_token")
+        # Also add `tokenizer.start_token` and `tokenizer.end_token` for
+        # compatibility with other tokenizers.
+        self._add_special_token("[CLS]", "start_token")
+        self._add_special_token("[SEP]", "end_token")
         super().__init__(proto=proto, **kwargs)
-
-    def set_proto(self, proto):
-        super().set_proto(proto)
-        if proto is not None:
-            for token in [
-                self.cls_token,
-                self.sep_token,
-                self.pad_token,
-                self.mask_token,
-            ]:
-                if token not in self.get_vocabulary():
-                    raise ValueError(
-                        f"Cannot find token `'{token}'` in the provided "
-                        f"`vocabulary`. Please provide `'{token}'` in your "
-                        "`vocabulary` or use a pretrained `vocabulary` name."
-                    )
-
-            self.cls_token_id = self.token_to_id(self.cls_token)
-            self.sep_token_id = self.token_to_id(self.sep_token)
-            self.pad_token_id = self.token_to_id(self.pad_token)
-            self.mask_token_id = self.token_to_id(self.mask_token)
-        else:
-            self.cls_token_id = None
-            self.sep_token_id = None
-            self.pad_token_id = None
-            self.mask_token_id = None

keras_hub/src/models/falcon/__init__.py

@@ -14,7 +14,6 @@
 
 from keras_hub.src.models.falcon.falcon_backbone import FalconBackbone
 from keras_hub.src.models.falcon.falcon_presets import backbone_presets
-from keras_hub.src.models.falcon.falcon_tokenizer import FalconTokenizer
 from keras_hub.src.utils.preset_utils import register_presets
 
-register_presets(backbone_presets, (FalconBackbone, FalconTokenizer))
+register_presets(backbone_presets, FalconBackbone)

keras_hub/src/models/falcon/falcon_causal_lm_preprocessor.py

@@ -12,19 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import keras
-from absl import logging
-
 from keras_hub.src.api_export import keras_hub_export
-from keras_hub.src.models.falcon.falcon_preprocessor import FalconPreprocessor
-from keras_hub.src.utils.keras_utils import (
-    convert_inputs_to_list_of_tensor_segments,
-)
-from keras_hub.src.utils.tensor_utils import strip_to_ragged
+from keras_hub.src.models.causal_lm_preprocessor import CausalLMPreprocessor
+from keras_hub.src.models.falcon.falcon_backbone import FalconBackbone
+from keras_hub.src.models.falcon.falcon_tokenizer import FalconTokenizer
 
 
 @keras_hub_export("keras_hub.models.FalconCausalLMPreprocessor")
-class FalconCausalLMPreprocessor(FalconPreprocessor):
+class FalconCausalLMPreprocessor(CausalLMPreprocessor):
     """Falcon Causal LM preprocessor.
 
     This preprocessing layer is meant for use with
@@ -91,83 +86,5 @@ class FalconCausalLMPreprocessor(FalconPreprocessor):
     ```
     """
 
-    def call(
-        self,
-        x,
-        y=None,
-        sample_weight=None,
-        sequence_length=None,
-    ):
-        if y is not None or sample_weight is not None:
-            logging.warning(
-                "`FalconCausalLMPreprocessor` generates `y` and `sample_weight` "
-                "based on your input data, but your data already contains `y` "
-                "or `sample_weight`. Your `y` and `sample_weight` will be "
-                "ignored."
-            )
-        sequence_length = sequence_length or self.sequence_length
-
-        x = convert_inputs_to_list_of_tensor_segments(x)[0]
-        x = self.tokenizer(x)
-        # Pad with one extra token to account for the truncation below.
-        token_ids, padding_mask = self.packer(
-            x,
-            sequence_length=sequence_length + 1,
-            add_start_value=self.add_start_token,
-            add_end_value=self.add_end_token,
-        )
-        # The last token does not have a next token, so we truncate it out.
-        x = {
-            "token_ids": token_ids[..., :-1],
-            "padding_mask": padding_mask[..., :-1],
-        }
-        # Target `y` will be the next token.
-        y, sample_weight = token_ids[..., 1:], padding_mask[..., 1:]
-        return keras.utils.pack_x_y_sample_weight(x, y, sample_weight)
-
-    def generate_preprocess(
-        self,
-        x,
-        sequence_length=None,
-    ):
-        """Convert strings to integer token input for generation.
-
-        Similar to calling the layer for training, this method takes in strings
-        or tensor strings, tokenizes and packs the input, and computes a padding
-        mask masking all inputs not filled in with a padded value.
-
-        Unlike calling the layer for training, this method does not compute
-        labels and will never append a `tokenizer.end_token_id` to the end of
-        the sequence (as generation is expected to continue at the end of the
-        inputted prompt).
-        """
-        if not self.built:
-            self.build(None)
-
-        x = convert_inputs_to_list_of_tensor_segments(x)[0]
-        x = self.tokenizer(x)
-        token_ids, padding_mask = self.packer(
-            x, sequence_length=sequence_length, add_end_value=False
-        )
-        return {
-            "token_ids": token_ids,
-            "padding_mask": padding_mask,
-        }
-
-    def generate_postprocess(
-        self,
-        x,
-    ):
-        """Convert integer token output to strings for generation.
-
-        This method reverses `generate_preprocess()`, by first removing all
-        padding and start/end tokens, and then converting the integer sequence
-        back to a string.
-        """
-        if not self.built:
-            self.build(None)
-
-        token_ids, padding_mask = x["token_ids"], x["padding_mask"]
-        ids_to_strip = (self.tokenizer.end_token_id,)
-        token_ids = strip_to_ragged(token_ids, padding_mask, ids_to_strip)
-        return self.tokenizer.detokenize(token_ids)
+    backbone_cls = FalconBackbone
+    tokenizer_cls = FalconTokenizer

keras_hub/src/models/falcon/falcon_preprocessor.py

@@ -17,11 +17,10 @@ import keras
 
 from keras_hub.src.api_export import keras_hub_export
 from keras_hub.src.layers.preprocessing.start_end_packer import StartEndPacker
+from keras_hub.src.models.falcon.falcon_backbone import FalconBackbone
 from keras_hub.src.models.falcon.falcon_tokenizer import FalconTokenizer
 from keras_hub.src.models.preprocessor import Preprocessor
-from keras_hub.src.utils.keras_utils import (
-    convert_inputs_to_list_of_tensor_segments,
-)
+from keras_hub.src.utils.tensor_utils import preprocessing_function
 
 
 @keras_hub_export("keras_hub.models.FalconPreprocessor")
@@ -107,6 +106,7 @@ class FalconPreprocessor(Preprocessor):
     ```
     """
 
+    backbone_cls = FalconBackbone
     tokenizer_cls = FalconTokenizer
 
     def __init__(
@@ -136,6 +136,7 @@ class FalconPreprocessor(Preprocessor):
         )
         self.built = True
 
+    @preprocessing_function
     def call(
         self,
         x,
@@ -143,17 +144,9 @@ class FalconPreprocessor(Preprocessor):
         sample_weight=None,
         sequence_length=None,
     ):
-        x = convert_inputs_to_list_of_tensor_segments(x)
-        if len(x) != 1:
-            raise ValueError(
-                "Falcon requires each input feature to contain only "
-                f"one segment, but received {len(x)}. If you are using Falcon "
-                "for a multi-segment classification task, please refer to "
-                "classification models like BERT or RoBERTa."
-            )
         sequence_length = sequence_length or self.sequence_length
         token_ids, padding_mask = self.packer(
-            self.tokenizer(x[0]),
+            self.tokenizer(x),
             sequence_length=sequence_length,
             add_start_value=self.add_start_token,
             add_end_value=self.add_end_token,

keras_hub/src/models/falcon/falcon_tokenizer.py

@@ -14,10 +14,16 @@
 
 
 from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.models.falcon.falcon_backbone import FalconBackbone
 from keras_hub.src.tokenizers.byte_pair_tokenizer import BytePairTokenizer
 
 
-@keras_hub_export("keras_hub.models.FalconTokenizer")
+@keras_hub_export(
+    [
+        "keras_hub.tokenizers.FalconTokenizer",
+        "keras_hub.models.FalconTokenizer",
+    ]
+)
 class FalconTokenizer(BytePairTokenizer):
     """Falcon tokenizer based on BytePairTokenizer.
 
@@ -27,8 +33,6 @@ class FalconTokenizer(BytePairTokenizer):
     models and provides a `from_preset()` method to automatically download
     a matching vocabulary for a Falcon preset.
 
-    This tokenizer does not provide truncation or padding of inputs.
-
     If input is a batch of strings (rank > 0), the layer will output a
     `tf.RaggedTensor` where the last dimension of the output is ragged.
 
@@ -65,46 +69,19 @@ class FalconTokenizer(BytePairTokenizer):
     ```
     """
 
+    backbone_cls = FalconBackbone
+
     def __init__(
         self,
         vocabulary=None,
         merges=None,
         **kwargs,
     ):
-        # Falcon uses the same start as end token, i.e., "<|endoftext|>".
-        self.end_token = self.start_token = "<|endoftext|>"
-
+        self._add_special_token("<|endoftext|>", "end_token")
+        self._add_special_token("<|endoftext|>", "start_token")
+        self.pad_token_id = 0
         super().__init__(
             vocabulary=vocabulary,
             merges=merges,
-            unsplittable_tokens=[self.end_token],
             **kwargs,
         )
-
-    def set_vocabulary_and_merges(self, vocabulary, merges):
-        super().set_vocabulary_and_merges(vocabulary, merges)
-
-        if vocabulary is not None:
-            # Check for necessary special tokens.
-            if self.end_token not in self.get_vocabulary():
-                raise ValueError(
-                    f"Cannot find token `'{self.end_token}'` in the provided "
-                    f"`vocabulary`. Please provide `'{self.end_token}'` in "
-                    "your `vocabulary` or use a pretrained `vocabulary` name."
-                )
-
-            self.end_token_id = self.token_to_id(self.end_token)
-            self.start_token_id = self.end_token_id
-            self.pad_token_id = 0
-        else:
-            self.end_token_id = None
-            self.start_token_id = None
-            self.pad_token_id = None
-
-    def get_config(self):
-        config = super().get_config()
-        # In the constructor, we pass the list of special tokens to the
-        # `unsplittable_tokens` arg of the superclass' constructor. Hence, we
-        # delete it from the config here.
-        del config["unsplittable_tokens"]
-        return config

keras_hub/src/models/gemma/__init__.py

@@ -14,7 +14,6 @@
 
 from keras_hub.src.models.gemma.gemma_backbone import GemmaBackbone
 from keras_hub.src.models.gemma.gemma_presets import backbone_presets
-from keras_hub.src.models.gemma.gemma_tokenizer import GemmaTokenizer
 from keras_hub.src.utils.preset_utils import register_presets
 
-register_presets(backbone_presets, (GemmaBackbone, GemmaTokenizer))
+register_presets(backbone_presets, GemmaBackbone)