def test_network_invocation(self): config = bert.BertPretrainerConfig( encoder=encoders.TransformerEncoderConfig(vocab_size=10, num_layers=1)) _ = bert.instantiate_bertpretrainer_from_cfg(config) # Invokes with classification heads. config = bert.BertPretrainerConfig( encoder=encoders.TransformerEncoderConfig(vocab_size=10, num_layers=1), cls_heads=[ bert.ClsHeadConfig(inner_dim=10, num_classes=2, name="next_sentence") ]) _ = bert.instantiate_bertpretrainer_from_cfg(config) with self.assertRaises(ValueError): config = bert.BertPretrainerConfig( encoder=encoders.TransformerEncoderConfig(vocab_size=10, num_layers=1), cls_heads=[ bert.ClsHeadConfig(inner_dim=10, num_classes=2, name="next_sentence"), bert.ClsHeadConfig(inner_dim=10, num_classes=2, name="next_sentence") ]) _ = bert.instantiate_bertpretrainer_from_cfg(config)
def build_model(self): if self._hub_module: encoder_from_hub = utils.get_encoder_from_hub(self._hub_module) return bert.instantiate_bertpretrainer_from_cfg( self.task_config.model, encoder_network=encoder_from_hub) else: return bert.instantiate_bertpretrainer_from_cfg(self.task_config.model)
def test_task(self): # Saves a checkpoint. pretrain_cfg = bert.BertPretrainerConfig(encoder=self._encoder_config, num_masked_tokens=20, cls_heads=[ bert.ClsHeadConfig( inner_dim=10, num_classes=3, name="next_sentence") ]) pretrain_model = bert.instantiate_bertpretrainer_from_cfg(pretrain_cfg) ckpt = tf.train.Checkpoint(model=pretrain_model, **pretrain_model.checkpoint_items) saved_path = ckpt.save(self.get_temp_dir()) config = question_answering.QuestionAnsweringConfig( init_checkpoint=saved_path, network=self._encoder_config, train_data=self._train_data_config) task = question_answering.QuestionAnsweringTask(config) model = task.build_model() metrics = task.build_metrics() dataset = task.build_inputs(config.train_data) iterator = iter(dataset) optimizer = tf.keras.optimizers.SGD(lr=0.1) task.train_step(next(iterator), model, optimizer, metrics=metrics) task.validation_step(next(iterator), model, metrics=metrics) task.initialize(model)
def test_task(self): config = sentence_prediction.SentencePredictionConfig( init_checkpoint=self.get_temp_dir(), model=self.get_model_config(2), train_data=self._train_data_config) task = sentence_prediction.SentencePredictionTask(config) model = task.build_model() metrics = task.build_metrics() dataset = task.build_inputs(config.train_data) iterator = iter(dataset) optimizer = tf.keras.optimizers.SGD(lr=0.1) task.train_step(next(iterator), model, optimizer, metrics=metrics) task.validation_step(next(iterator), model, metrics=metrics) # Saves a checkpoint. pretrain_cfg = bert.BertPretrainerConfig( encoder=encoders.TransformerEncoderConfig(vocab_size=30522, num_layers=1), num_masked_tokens=20, cls_heads=[ bert.ClsHeadConfig(inner_dim=10, num_classes=3, name="next_sentence") ]) pretrain_model = bert.instantiate_bertpretrainer_from_cfg(pretrain_cfg) ckpt = tf.train.Checkpoint(model=pretrain_model, **pretrain_model.checkpoint_items) ckpt.save(config.init_checkpoint) task.initialize(model)
def test_checkpoint_items(self): config = bert.BertPretrainerConfig( encoder=encoders.TransformerEncoderConfig(vocab_size=10, num_layers=1), cls_heads=[ bert.ClsHeadConfig(inner_dim=10, num_classes=2, name="next_sentence") ]) encoder = bert.instantiate_bertpretrainer_from_cfg(config) self.assertSameElements(encoder.checkpoint_items.keys(), ["encoder", "next_sentence.pooler_dense"])
def build_model(self): if self._hub_module: input_word_ids = tf.keras.layers.Input(shape=(None, ), dtype=tf.int32, name='input_word_ids') input_mask = tf.keras.layers.Input(shape=(None, ), dtype=tf.int32, name='input_mask') input_type_ids = tf.keras.layers.Input(shape=(None, ), dtype=tf.int32, name='input_type_ids') bert_model = hub.KerasLayer(self._hub_module, trainable=True) pooled_output, sequence_output = bert_model( [input_word_ids, input_mask, input_type_ids]) encoder_from_hub = tf.keras.Model( inputs=[input_word_ids, input_mask, input_type_ids], outputs=[sequence_output, pooled_output]) return bert.instantiate_bertpretrainer_from_cfg( self.task_config.network, encoder_network=encoder_from_hub) else: return bert.instantiate_bertpretrainer_from_cfg( self.task_config.network)
def test_task(self, version_2_with_negative, tokenization): # Saves a checkpoint. pretrain_cfg = bert.BertPretrainerConfig(encoder=self._encoder_config, num_masked_tokens=20, cls_heads=[ bert.ClsHeadConfig( inner_dim=10, num_classes=3, name="next_sentence") ]) pretrain_model = bert.instantiate_bertpretrainer_from_cfg(pretrain_cfg) ckpt = tf.train.Checkpoint(model=pretrain_model, **pretrain_model.checkpoint_items) saved_path = ckpt.save(self.get_temp_dir()) config = question_answering.QuestionAnsweringConfig( init_checkpoint=saved_path, model=self._encoder_config, train_data=self._train_data_config, validation_data=self._get_validation_data_config( version_2_with_negative)) self._run_task(config)
def build_model(self): return bert.instantiate_bertpretrainer_from_cfg(self.task_config.model)