| 182 | |
| 183 | @keras_serializable |
| 184 | class TFCTRLMainLayer(tf.keras.layers.Layer): |
| 185 | config_class = CTRLConfig |
| 186 | |
| 187 | def __init__(self, config, **kwargs): |
| 188 | super().__init__(**kwargs) |
| 189 | self.output_hidden_states = config.output_hidden_states |
| 190 | self.output_attentions = config.output_attentions |
| 191 | self.use_cache = config.use_cache |
| 192 | |
| 193 | self.d_model_size = config.n_embd |
| 194 | self.num_layers = config.n_layer |
| 195 | |
| 196 | self.pos_encoding = positional_encoding(config.n_positions, self.d_model_size) |
| 197 | |
| 198 | self.w = TFSharedEmbeddings( |
| 199 | config.vocab_size, config.n_embd, initializer_range=config.initializer_range, name="w" |
| 200 | ) |
| 201 | |
| 202 | self.dropout = tf.keras.layers.Dropout(config.embd_pdrop) |
| 203 | self.h = [ |
| 204 | TFEncoderLayer( |
| 205 | config.n_embd, |
| 206 | config.n_head, |
| 207 | config.dff, |
| 208 | config.resid_pdrop, |
| 209 | config.layer_norm_epsilon, |
| 210 | name="h_._{}".format(i), |
| 211 | ) |
| 212 | for i in range(config.n_layer) |
| 213 | ] |
| 214 | self.layernorm = tf.keras.layers.LayerNormalization(epsilon=config.layer_norm_epsilon, name="layernorm") |
| 215 | |
| 216 | def get_input_embeddings(self): |
| 217 | return self.w |
| 218 | |
| 219 | def set_input_embeddings(self, value): |
| 220 | self.w.weight = value |
| 221 | self.w.vocab_size = value.shape[0] |
| 222 | |
| 223 | def _resize_token_embeddings(self, new_num_tokens): |
| 224 | raise NotImplementedError |
| 225 | |
| 226 | def _prune_heads(self, heads_to_prune): |
| 227 | """ Prunes heads of the model. |
| 228 | heads_to_prune: dict of {layer_num: list of heads to prune in this layer} |
| 229 | """ |
| 230 | raise NotImplementedError |
| 231 | |
| 232 | def call( |
| 233 | self, |
| 234 | inputs, |
| 235 | past=None, |
| 236 | attention_mask=None, |
| 237 | token_type_ids=None, |
| 238 | position_ids=None, |
| 239 | head_mask=None, |
| 240 | inputs_embeds=None, |
| 241 | use_cache=None, |