keras-hub-nightly 0.15.0.dev20240911134614__py3-none-any.whl → 0.16.0.dev20240915160609__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (36) hide show
  1. keras_hub/api/models/__init__.py +22 -17
  2. keras_hub/src/models/albert/albert_text_classifier.py +6 -1
  3. keras_hub/src/models/bert/bert_text_classifier.py +6 -1
  4. keras_hub/src/models/deberta_v3/deberta_v3_text_classifier.py +6 -1
  5. keras_hub/src/models/distil_bert/distil_bert_text_classifier.py +6 -1
  6. keras_hub/src/models/f_net/f_net_text_classifier.py +6 -1
  7. keras_hub/src/models/gpt2/gpt2_preprocessor.py +7 -78
  8. keras_hub/src/models/pali_gemma/pali_gemma_tokenizer.py +1 -1
  9. keras_hub/src/models/preprocessor.py +1 -5
  10. keras_hub/src/models/resnet/resnet_backbone.py +3 -16
  11. keras_hub/src/models/resnet/resnet_image_classifier.py +17 -0
  12. keras_hub/src/models/resnet/resnet_presets.py +12 -12
  13. keras_hub/src/models/roberta/roberta_text_classifier.py +6 -1
  14. keras_hub/src/models/task.py +6 -6
  15. keras_hub/src/models/text_classifier.py +12 -1
  16. keras_hub/src/models/xlm_roberta/xlm_roberta_text_classifier.py +6 -1
  17. keras_hub/src/tests/test_case.py +13 -0
  18. keras_hub/src/utils/preset_utils.py +14 -32
  19. keras_hub/src/utils/timm/convert_resnet.py +0 -1
  20. keras_hub/src/utils/timm/preset_loader.py +6 -7
  21. keras_hub/src/version_utils.py +1 -1
  22. {keras_hub_nightly-0.15.0.dev20240911134614.dist-info → keras_hub_nightly-0.16.0.dev20240915160609.dist-info}/METADATA +1 -1
  23. {keras_hub_nightly-0.15.0.dev20240911134614.dist-info → keras_hub_nightly-0.16.0.dev20240915160609.dist-info}/RECORD +25 -36
  24. {keras_hub_nightly-0.15.0.dev20240911134614.dist-info → keras_hub_nightly-0.16.0.dev20240915160609.dist-info}/WHEEL +1 -1
  25. keras_hub/src/models/bart/bart_preprocessor.py +0 -264
  26. keras_hub/src/models/bloom/bloom_preprocessor.py +0 -178
  27. keras_hub/src/models/electra/electra_preprocessor.py +0 -155
  28. keras_hub/src/models/falcon/falcon_preprocessor.py +0 -180
  29. keras_hub/src/models/gemma/gemma_preprocessor.py +0 -184
  30. keras_hub/src/models/gpt_neo_x/gpt_neo_x_preprocessor.py +0 -138
  31. keras_hub/src/models/llama/llama_preprocessor.py +0 -182
  32. keras_hub/src/models/llama3/llama3_preprocessor.py +0 -23
  33. keras_hub/src/models/mistral/mistral_preprocessor.py +0 -183
  34. keras_hub/src/models/opt/opt_preprocessor.py +0 -181
  35. keras_hub/src/models/phi3/phi3_preprocessor.py +0 -183
  36. {keras_hub_nightly-0.15.0.dev20240911134614.dist-info → keras_hub_nightly-0.16.0.dev20240915160609.dist-info}/top_level.txt +0 -0
@@ -1,23 +0,0 @@
1
- # Copyright 2024 The KerasHub Authors
2
- #
3
- # Licensed under the Apache License, Version 2.0 (the "License");
4
- # you may not use this file except in compliance with the License.
5
- # You may obtain a copy of the License at
6
- #
7
- # https://www.apache.org/licenses/LICENSE-2.0
8
- #
9
- # Unless required by applicable law or agreed to in writing, software
10
- # distributed under the License is distributed on an "AS IS" BASIS,
11
- # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
- # See the License for the specific language governing permissions and
13
- # limitations under the License.
14
- from keras_hub.src.api_export import keras_hub_export
15
- from keras_hub.src.models.llama3.llama3_backbone import Llama3Backbone
16
- from keras_hub.src.models.llama3.llama3_tokenizer import Llama3Tokenizer
17
- from keras_hub.src.models.llama.llama_preprocessor import LlamaPreprocessor
18
-
19
-
20
- @keras_hub_export("keras_hub.models.Llama3Preprocessor")
21
- class Llama3Preprocessor(LlamaPreprocessor):
22
- backbone_cls = Llama3Backbone
23
- tokenizer_cls = Llama3Tokenizer
@@ -1,183 +0,0 @@
1
- # Copyright 2024 The KerasHub Authors
2
- #
3
- # Licensed under the Apache License, Version 2.0 (the "License");
4
- # you may not use this file except in compliance with the License.
5
- # You may obtain a copy of the License at
6
- #
7
- # https://www.apache.org/licenses/LICENSE-2.0
8
- #
9
- # Unless required by applicable law or agreed to in writing, software
10
- # distributed under the License is distributed on an "AS IS" BASIS,
11
- # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
- # See the License for the specific language governing permissions and
13
- # limitations under the License.
14
-
15
- import keras
16
-
17
- from keras_hub.src.api_export import keras_hub_export
18
- from keras_hub.src.layers.preprocessing.start_end_packer import StartEndPacker
19
- from keras_hub.src.models.mistral.mistral_backbone import MistralBackbone
20
- from keras_hub.src.models.mistral.mistral_tokenizer import MistralTokenizer
21
- from keras_hub.src.models.preprocessor import Preprocessor
22
- from keras_hub.src.utils.tensor_utils import preprocessing_function
23
-
24
-
25
- @keras_hub_export("keras_hub.models.MistralPreprocessor")
26
- class MistralPreprocessor(Preprocessor):
27
- """A Mistral preprocessing layer which tokenizes and packs inputs.
28
-
29
- This preprocessing layer will do three things:
30
-
31
- 1. Tokenize any number of input segments using the `tokenizer`.
32
- 2. Pack the inputs together using a `keras_hub.layers.StartEndPacker`.
33
- with the appropriate tokens.
34
- 3. Construct a dictionary with keys `"token_ids"`, and `"padding_mask"`
35
- that can be passed directly to `keras_hub.models.MistralBackbone`.
36
-
37
- This layer can be used directly with `tf.data.Dataset.map` to preprocess
38
- string data in the `(x, y, sample_weight)` format used by
39
- `keras.Model.fit`.
40
-
41
- Args:
42
- tokenizer: A `keras_hub.models.MistralTokenizer` instance.
43
- sequence_length: The length of the packed inputs.
44
- add_start_token: If `True`, the preprocessor will prepend the tokenizer
45
- start token to each input sequence. Default is `True`.
46
- add_end_token: If `True`, the preprocessor will append the tokenizer
47
- end token to each input sequence. Default is `False`.
48
-
49
- Call arguments:
50
- x: A tensor of single string sequences, or a tuple of multiple
51
- tensor sequences to be packed together. Inputs may be batched or
52
- unbatched. For single sequences, raw python inputs will be converted
53
- to tensors. For multiple sequences, pass tensors directly.
54
- y: Any label data. Will be passed through unaltered.
55
- sample_weight: Any label weight data. Will be passed through unaltered.
56
- sequence_length: Pass to override the configured `sequence_length` of
57
- the layer.
58
-
59
- Examples:
60
-
61
- Directly calling the from_preset().
62
- ```python
63
- preprocessor = keras_hub.models.MistralPreprocessor.from_preset(
64
- "mistral_base_en"
65
- )
66
-
67
- # Tokenize and pack a single sentence.
68
- preprocessor("The quick brown fox jumped.")
69
-
70
- # Tokenize and a batch of single sentences.
71
- preprocessor(["The quick brown fox jumped.", "Call me Ishmael."])
72
-
73
- # Preprocess a batch of sentence pairs.
74
- # When handling multiple sequences, always convert to tensors first!
75
- first = tf.constant(["The quick brown fox jumped.", "Call me Ishmael."])
76
- second = tf.constant(["The fox tripped.", "Oh look, a whale."])
77
- preprocessor((first, second))
78
- ```
79
-
80
- Mapping with `tf.data.Dataset`.
81
- ```python
82
- preprocessor = keras_hub.models.MistralPreprocessor.from_preset(
83
- "mistral_base_en"
84
- )
85
- first = tf.constant(["The quick brown fox jumped.", "Call me Ishmael."])
86
- second = tf.constant(["The fox tripped.", "Oh look, a whale."])
87
- label = tf.constant([1, 1])
88
-
89
- # Map labeled single sentences.
90
- ds = tf.data.Dataset.from_tensor_slices((first, label))
91
- ds = ds.map(preprocessor, num_parallel_calls=tf.data.AUTOTUNE)
92
-
93
- # Map unlabeled single sentences.
94
- ds = tf.data.Dataset.from_tensor_slices(first)
95
- ds = ds.map(preprocessor, num_parallel_calls=tf.data.AUTOTUNE)
96
-
97
- # Map labeled sentence pairs.
98
- ds = tf.data.Dataset.from_tensor_slices(((first, second), label))
99
- ds = ds.map(preprocessor, num_parallel_calls=tf.data.AUTOTUNE)
100
-
101
- # Map unlabeled sentence pairs.
102
- ds = tf.data.Dataset.from_tensor_slices((first, second))
103
-
104
- # Watch out for tf.data's default unpacking of tuples here!
105
- # Best to invoke the `preprocessor` directly in this case.
106
- ds = ds.map(
107
- lambda first, second: preprocessor(x=(first, second)),
108
- num_parallel_calls=tf.data.AUTOTUNE,
109
- )
110
- ```
111
- """
112
-
113
- backbone_cls = MistralBackbone
114
- tokenizer_cls = MistralTokenizer
115
-
116
- def __init__(
117
- self,
118
- tokenizer,
119
- sequence_length=1024,
120
- add_start_token=True,
121
- add_end_token=False,
122
- **kwargs,
123
- ):
124
- super().__init__(**kwargs)
125
- self.tokenizer = tokenizer
126
- self.packer = None
127
- self.add_start_token = add_start_token
128
- self.add_end_token = add_end_token
129
- self.sequence_length = sequence_length
130
-
131
- def build(self, input_shape):
132
- # Defer packer creation to `build()` so that we can be sure tokenizer
133
- # assets have loaded when restoring a saved model.
134
- self.packer = StartEndPacker(
135
- start_value=self.tokenizer.start_token_id,
136
- end_value=self.tokenizer.end_token_id,
137
- sequence_length=self.sequence_length,
138
- return_padding_mask=True,
139
- )
140
- self.built = True
141
-
142
- def get_config(self):
143
- config = super().get_config()
144
- config.update(
145
- {
146
- "sequence_length": self.sequence_length,
147
- "add_start_token": self.add_start_token,
148
- "add_end_token": self.add_end_token,
149
- }
150
- )
151
- return config
152
-
153
- @preprocessing_function
154
- def call(
155
- self,
156
- x,
157
- y=None,
158
- sample_weight=None,
159
- sequence_length=None,
160
- ):
161
- sequence_length = sequence_length or self.sequence_length
162
- token_ids, padding_mask = self.packer(
163
- self.tokenizer(x),
164
- sequence_length=sequence_length,
165
- add_start_value=self.add_start_token,
166
- add_end_value=self.add_end_token,
167
- )
168
- x = {
169
- "token_ids": token_ids,
170
- "padding_mask": padding_mask,
171
- }
172
- return keras.utils.pack_x_y_sample_weight(x, y, sample_weight)
173
-
174
- @property
175
- def sequence_length(self):
176
- """The padded length of model input sequences."""
177
- return self._sequence_length
178
-
179
- @sequence_length.setter
180
- def sequence_length(self, value):
181
- self._sequence_length = value
182
- if self.packer is not None:
183
- self.packer.sequence_length = value
@@ -1,181 +0,0 @@
1
- # Copyright 2024 The KerasHub Authors
2
- #
3
- # Licensed under the Apache License, Version 2.0 (the "License");
4
- # you may not use this file except in compliance with the License.
5
- # You may obtain a copy of the License at
6
- #
7
- # https://www.apache.org/licenses/LICENSE-2.0
8
- #
9
- # Unless required by applicable law or agreed to in writing, software
10
- # distributed under the License is distributed on an "AS IS" BASIS,
11
- # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
- # See the License for the specific language governing permissions and
13
- # limitations under the License.
14
-
15
-
16
- import keras
17
-
18
- from keras_hub.src.api_export import keras_hub_export
19
- from keras_hub.src.layers.preprocessing.start_end_packer import StartEndPacker
20
- from keras_hub.src.models.opt.opt_backbone import OPTBackbone
21
- from keras_hub.src.models.opt.opt_tokenizer import OPTTokenizer
22
- from keras_hub.src.models.preprocessor import Preprocessor
23
- from keras_hub.src.utils.tensor_utils import preprocessing_function
24
-
25
-
26
- @keras_hub_export("keras_hub.models.OPTPreprocessor")
27
- class OPTPreprocessor(Preprocessor):
28
- """OPT preprocessing layer which tokenizes and packs inputs.
29
-
30
- This preprocessing layer will do 2 things:
31
-
32
- - Tokenize the input using the `tokenizer`.
33
- - Construct a dictionary with keys `"token_ids"`, `"padding_mask"`, that can
34
- be passed directly to a `keras_hub.models.OPTBackbone`.
35
-
36
- This layer can be used directly with `tf.data.Dataset.map` to preprocess
37
- string data in the `(x, y, sample_weight)` format used by
38
- `keras.Model.fit`.
39
-
40
- The call method of this layer accepts three arguments, `x`, `y`, and
41
- `sample_weight`. `x` can be a python string or tensor representing a single
42
- segment, a list of python strings representing a batch of single segments,
43
- or a list of tensors representing multiple segments to be packed together.
44
- `y` and `sample_weight` are both optional, can have any format, and will be
45
- passed through unaltered.
46
-
47
- `OPTPreprocessor` forces the input to have only one segment, as OPT is
48
- mainly used for generation tasks. For tasks having multi-segment inputs
49
- like "glue/mnli", please use a model designed for classification purposes
50
- such as BERT or RoBERTa.
51
-
52
- Args:
53
- tokenizer: A `keras_hub.models.OPTTokenizer` instance.
54
- sequence_length: The length of the packed inputs.
55
- add_start_token: If `True`, the preprocessor will append the tokenizer
56
- start token to each input sequence.
57
- add_end_token: If `True`, the preprocessor will append the tokenizer
58
- end token to each input sequence.
59
-
60
- Call arguments:
61
- x: A string, `tf.Tensor` or list of python strings.
62
- y: Any label data. Will be passed through unaltered.
63
- sample_weight: Any label weight data. Will be passed through unaltered.
64
- sequence_length: Pass to override the configured `sequence_length` of
65
- the layer.
66
-
67
- Examples:
68
-
69
- Directly calling the layer on data.
70
- ```python
71
- preprocessor = keras_hub.models.OPTPreprocessor.from_preset("opt_125m_en")
72
-
73
- # Tokenize and pack a single sentence.
74
- preprocessor("The quick brown fox jumped.")
75
-
76
- # Tokenize a batch of single sentences.
77
- preprocessor(["The quick brown fox jumped.", "Call me Ishmael."])
78
-
79
- # Custom vocabulary.
80
- features = ["a quick fox.", "a fox quick."]
81
- vocab = {"<|endoftext|>": 0, "a": 4, "Ġquick": 5, "Ġfox": 6}
82
- merges = ["Ġ q", "u i", "c k", "ui ck", "Ġq uick"]
83
- merges += ["Ġ f", "o x", "Ġf ox"]
84
- tokenizer = keras_hub.models.OPTTokenizer(
85
- vocabulary=vocab,
86
- merges=merges,
87
- )
88
- preprocessor = keras_hub.models.OPTPreprocessor(tokenizer=tokenizer)
89
- preprocessor("The quick brown fox jumped.")
90
- ```
91
-
92
- Mapping with `tf.data.Dataset`.
93
- ```python
94
- preprocessor = keras_hub.models.OPTPreprocessor.from_preset("opt_125m_en")
95
-
96
- text = tf.constant(["The quick brown fox jumped.", "Call me Ishmael."])
97
- label = tf.constant([1, 1])
98
-
99
- # Map labeled single sentences.
100
- ds = tf.data.Dataset.from_tensor_slices((text, label))
101
- ds = ds.map(preprocessor, num_parallel_calls=tf.data.AUTOTUNE)
102
-
103
- # Map unlabeled single sentences.
104
- ds = tf.data.Dataset.from_tensor_slices(text)
105
- ds = ds.map(preprocessor, num_parallel_calls=tf.data.AUTOTUNE)
106
- ```
107
- """
108
-
109
- backbone_cls = OPTBackbone
110
- tokenizer_cls = OPTTokenizer
111
-
112
- def __init__(
113
- self,
114
- tokenizer,
115
- sequence_length=1024,
116
- add_start_token=True,
117
- add_end_token=True,
118
- **kwargs,
119
- ):
120
- super().__init__(**kwargs)
121
-
122
- self.tokenizer = tokenizer
123
- self.packer = None
124
- self.sequence_length = sequence_length
125
- self.add_start_token = add_start_token
126
- self.add_end_token = add_end_token
127
-
128
- def build(self, input_shape):
129
- # Defer packer creation to `build()` so that we can be sure tokenizer
130
- # assets have loaded when restoring a saved model.
131
- self.packer = StartEndPacker(
132
- start_value=self.tokenizer.start_token_id,
133
- end_value=self.tokenizer.end_token_id,
134
- pad_value=self.tokenizer.pad_token_id,
135
- sequence_length=self.sequence_length,
136
- return_padding_mask=True,
137
- )
138
- self.built = True
139
-
140
- def get_config(self):
141
- config = super().get_config()
142
- config.update(
143
- {
144
- "sequence_length": self.sequence_length,
145
- "add_start_token": self.add_start_token,
146
- "add_end_token": self.add_end_token,
147
- }
148
- )
149
- return config
150
-
151
- @preprocessing_function
152
- def call(
153
- self,
154
- x,
155
- y=None,
156
- sample_weight=None,
157
- sequence_length=None,
158
- ):
159
- sequence_length = sequence_length or self.sequence_length
160
- token_ids, padding_mask = self.packer(
161
- self.tokenizer(x),
162
- sequence_length=sequence_length,
163
- add_start_value=self.add_start_token,
164
- add_end_value=self.add_end_token,
165
- )
166
- x = {
167
- "token_ids": token_ids,
168
- "padding_mask": padding_mask,
169
- }
170
- return keras.utils.pack_x_y_sample_weight(x, y, sample_weight)
171
-
172
- @property
173
- def sequence_length(self):
174
- """The padded length of model input sequences."""
175
- return self._sequence_length
176
-
177
- @sequence_length.setter
178
- def sequence_length(self, value):
179
- self._sequence_length = value
180
- if self.packer is not None:
181
- self.packer.sequence_length = value
@@ -1,183 +0,0 @@
1
- # Copyright 2024 The KerasHub Authors
2
- #
3
- # Licensed under the Apache License, Version 2.0 (the "License");
4
- # you may not use this file except in compliance with the License.
5
- # You may obtain a copy of the License at
6
- #
7
- # https://www.apache.org/licenses/LICENSE-2.0
8
- #
9
- # Unless required by applicable law or agreed to in writing, software
10
- # distributed under the License is distributed on an "AS IS" BASIS,
11
- # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
- # See the License for the specific language governing permissions and
13
- # limitations under the License.
14
- import keras
15
-
16
- from keras_hub.src.api_export import keras_hub_export
17
- from keras_hub.src.layers.preprocessing.start_end_packer import StartEndPacker
18
- from keras_hub.src.models.phi3.phi3_backbone import Phi3Backbone
19
- from keras_hub.src.models.phi3.phi3_tokenizer import Phi3Tokenizer
20
- from keras_hub.src.models.preprocessor import Preprocessor
21
- from keras_hub.src.utils.tensor_utils import preprocessing_function
22
-
23
-
24
- @keras_hub_export("keras_hub.models.Phi3Preprocessor")
25
- class Phi3Preprocessor(Preprocessor):
26
- """A Phi3 preprocessing layer which tokenizes and packs inputs.
27
-
28
- This preprocessing layer will do three things:
29
-
30
- 1. Tokenize any number of input segments using the `tokenizer`.
31
- 2. Pack the inputs together using a `keras_hub.layers.StartEndPacker`.
32
- with the appropriate tokens.
33
- 3. Construct a dictionary with keys `"token_ids"`, and `"padding_mask"`
34
- that can be passed directly to `keras_hub.models.Phi3Backbone`.
35
-
36
- This layer can be used directly with `tf.data.Dataset.map` to preprocess
37
- string data in the `(x, y, sample_weight)` format used by
38
- `keras.Model.fit`.
39
-
40
- Args:
41
- tokenizer: A `keras_hub.models.Phi3Tokenizer` instance.
42
- sequence_length: The length of the packed inputs.
43
- add_start_token: If `True`, the preprocessor will prepend the tokenizer
44
- start token to each input sequence. Default is `True`.
45
- add_end_token: If `True`, the preprocessor will append the tokenizer
46
- end token to each input sequence. Default is `False`.
47
-
48
- Call arguments:
49
- x: A tensor of single string sequences, or a tuple of multiple
50
- tensor sequences to be packed together. Inputs may be batched or
51
- unbatched. For single sequences, raw python inputs will be converted
52
- to tensors. For multiple sequences, pass tensors directly.
53
- y: Any label data. Will be passed through unaltered.
54
- sample_weight: Any label weight data. Will be passed through unaltered.
55
- sequence_length: Pass to override the configured `sequence_length` of
56
- the layer.
57
-
58
- Examples:
59
-
60
- Directly calling the from_preset().
61
- ```python
62
- preprocessor = keras_hub.models.Phi3Preprocessor.from_preset(
63
- ""
64
- )
65
-
66
- # Tokenize and pack a single sentence.
67
- preprocessor("The quick brown fox jumped.")
68
-
69
- # Tokenize and a batch of single sentences.
70
- preprocessor(["The quick brown fox jumped.", "Call me Ishmael."])
71
-
72
- # Preprocess a batch of sentence pairs.
73
- # When handling multiple sequences, always convert to tensors first!
74
- first = tf.constant(["The quick brown fox jumped.", "Call me Ishmael."])
75
- second = tf.constant(["The fox tripped.", "Oh look, a whale."])
76
- preprocessor((first, second))
77
- ```
78
-
79
- Mapping with `tf.data.Dataset`.
80
- ```python
81
- preprocessor = keras_hub.models.Phi3Preprocessor.from_preset(
82
- "phi3_mini_4k_instruct_en"
83
- )
84
- first = tf.constant(["The quick brown fox jumped.", "Call me Ishmael."])
85
- second = tf.constant(["The fox tripped.", "Oh look, a whale."])
86
- label = tf.constant([1, 1])
87
-
88
- # Map labeled single sentences.
89
- ds = tf.data.Dataset.from_tensor_slices((first, label))
90
- ds = ds.map(preprocessor, num_parallel_calls=tf.data.AUTOTUNE)
91
-
92
- # Map unlabeled single sentences.
93
- ds = tf.data.Dataset.from_tensor_slices(first)
94
- ds = ds.map(preprocessor, num_parallel_calls=tf.data.AUTOTUNE)
95
-
96
- # Map labeled sentence pairs.
97
- ds = tf.data.Dataset.from_tensor_slices(((first, second), label))
98
- ds = ds.map(preprocessor, num_parallel_calls=tf.data.AUTOTUNE)
99
-
100
- # Map unlabeled sentence pairs.
101
- ds = tf.data.Dataset.from_tensor_slices((first, second))
102
-
103
- # Watch out for tf.data's default unpacking of tuples here!
104
- # Best to invoke the `preprocessor` directly in this case.
105
- ds = ds.map(
106
- lambda first, second: preprocessor(x=(first, second)),
107
- num_parallel_calls=tf.data.AUTOTUNE,
108
- )
109
- ```
110
- """
111
-
112
- backbone_cls = Phi3Backbone
113
- tokenizer_cls = Phi3Tokenizer
114
-
115
- def __init__(
116
- self,
117
- tokenizer,
118
- sequence_length=1024,
119
- add_start_token=True,
120
- add_end_token=False,
121
- **kwargs,
122
- ):
123
- super().__init__(**kwargs)
124
- self.tokenizer = tokenizer
125
- self.packer = None
126
- self.add_start_token = add_start_token
127
- self.add_end_token = add_end_token
128
- self.sequence_length = sequence_length
129
-
130
- def build(self, input_shape):
131
- # Defer packer creation to `build()` so that we can be sure tokenizer
132
- # assets have loaded when restoring a saved model.
133
- self.packer = StartEndPacker(
134
- start_value=self.tokenizer.start_token_id,
135
- end_value=self.tokenizer.end_token_id,
136
- pad_value=self.tokenizer.pad_token_id,
137
- sequence_length=self.sequence_length,
138
- return_padding_mask=True,
139
- )
140
- self.built = True
141
-
142
- def get_config(self):
143
- config = super().get_config()
144
- config.update(
145
- {
146
- "sequence_length": self.sequence_length,
147
- "add_start_token": self.add_start_token,
148
- "add_end_token": self.add_end_token,
149
- }
150
- )
151
- return config
152
-
153
- @preprocessing_function
154
- def call(
155
- self,
156
- x,
157
- y=None,
158
- sample_weight=None,
159
- sequence_length=None,
160
- ):
161
- sequence_length = sequence_length or self.sequence_length
162
- token_ids, padding_mask = self.packer(
163
- self.tokenizer(x),
164
- sequence_length=sequence_length,
165
- add_start_value=self.add_start_token,
166
- add_end_value=self.add_end_token,
167
- )
168
- x = {
169
- "token_ids": token_ids,
170
- "padding_mask": padding_mask,
171
- }
172
- return keras.utils.pack_x_y_sample_weight(x, y, sample_weight)
173
-
174
- @property
175
- def sequence_length(self):
176
- """The padded length of model input sequences."""
177
- return self._sequence_length
178
-
179
- @sequence_length.setter
180
- def sequence_length(self, value):
181
- self._sequence_length = value
182
- if self.packer is not None:
183
- self.packer.sequence_length = value