def create_encoder_model(self, is_training, input_ids, input_mask,
segment_ids):
return canine_modeling.CanineModel(config=self.model_config,
atom_input_ids=input_ids,
atom_input_mask=input_mask,
atom_segment_ids=segment_ids,
is_training=is_training)
def create_model(self, seed=None):
input_ids = ids_tensor([self.batch_size, self.seq_length],
self.vocab_size,
seed=seed,
dynamic_batch_size=self.dynamic_batch_size)
if seed is not None:
seed *= 7
input_mask = ids_tensor([self.batch_size, self.seq_length],
vocab_size=2,
seed=seed,
dynamic_batch_size=self.dynamic_batch_size)
if seed is not None:
seed *= 5
token_type_ids = ids_tensor(
[self.batch_size, self.seq_length],
self.type_vocab_size,
seed=seed,
dynamic_batch_size=self.dynamic_batch_size)
config = modeling.CanineModelConfig(
hidden_size=self.hidden_size,
num_hidden_layers=self.num_hidden_layers,
num_attention_heads=self.num_attention_heads,
intermediate_size=self.intermediate_size,
max_positions=self.max_position_embeddings,
downsampling_rate=self.downsampling_rate)
model = modeling.CanineModel(config=config,
atom_input_ids=input_ids,
atom_input_mask=input_mask,
atom_segment_ids=token_type_ids,
is_training=self.is_training)
outputs = {
"pooled_output": model.get_pooled_output(),
"sequence_output": model.get_sequence_output(),
"downsampled_layers": model.get_downsampled_layers(),
}
return outputs