def __call__(self, input_ids, input_mask, type_ids, masked_lm_positions, masked_lm_labels, masked_lm_weights, next_sentence_labels, deterministic=False): """Applies pre-training model on inputs. Args: input_ids: <int>[BATCH_SIZE, MAX_SEQ_LENGTH] tokenized inputs. input_mask: <bool>[BATCH_SIZE, MAX_SEQ_LENGTH] mask separating actual inputs from padding. Only used by BERT. type_ids: <int>[BATCH_SIZE, MAX_SEQ_LENGTH] Ids partitioning input into different types. masked_lm_positions: <int>[BATCH_SIZE, MAX_PREDICTIONS_PER_SEQ] indices indicating which inputs are masked. masked_lm_labels: <int>[BATCH_SIZE, MAX_PREDICTIONS_PER_SEQ] true labels for masked inputs. masked_lm_weights: <float>[BATCH_SIZE, MAX_PREDICTIONS_PER_SEQ] relative weighting for masked inputs. next_sentence_labels: <int>[BATCH_SIZE, 1] Labels for next sentence prediction task. deterministic: Whether or not to apply dropout to input. Returns: Loss and metrics for given inputs. """ sequence_output, pooled_output = EncoderModel( self.config, random_seed=self.random_seed, name="encoder")(input_ids, input_mask, type_ids, deterministic=deterministic) masked_lm_output = layers.gather(sequence_output, masked_lm_positions) masked_lm_output = nn.Dense(self.config.d_emb, kernel_init=default_kernel_init, name="predictions_dense")(masked_lm_output) masked_lm_output = nn.gelu(masked_lm_output) masked_lm_output = nn.LayerNorm( epsilon=LAYER_NORM_EPSILON, name="predictions_layer_norm")(masked_lm_output) masked_lm_logits = layers.OutputProjection( kernel=self._get_embedding_table(), name="predictions_output")(masked_lm_output) next_sentence_logits = layers.OutputProjection( n_out=2, kernel_init=default_kernel_init, name="classification")(pooled_output) return _compute_pretraining_metrics(masked_lm_logits, next_sentence_logits, masked_lm_labels, masked_lm_weights, next_sentence_labels)
def test_kernel_output_projection(self): batch_size = 2 max_seq_length = 14 hidden_dim = 8 rng = jax.random.PRNGKey(0) rng, kernel_rng = jax.random.split(rng) kernel = jax.random.uniform(kernel_rng, (NUM_CLASSES, hidden_dim), minval=-1.0, maxval=1.0) kernel_projection = layers.OutputProjection(kernel=jnp.asarray(kernel), name="predictions_output") init_batch = { "inputs": jnp.ones((1, max_seq_length, hidden_dim), jnp.float32) } params = init_layer_variables(rng, kernel_projection, init_batch)["params"] expected_keys = {"output_bias"} self.assertEqual(params.keys(), expected_keys) rng, init_rng = jax.random.split(rng) inputs = jax.random.randint(init_rng, (batch_size, max_seq_length, hidden_dim), minval=0, maxval=10) outputs = kernel_projection.apply({"params": params}, inputs=inputs) self.assertEqual(outputs.shape, (batch_size, max_seq_length, NUM_CLASSES))
def test_classification_output_projection(self): batch_size = 2 max_seq_length = 14 hidden_dim = 8 rng = jax.random.PRNGKey(0) classification_projection = layers.OutputProjection( n_out=NUM_CLASSES, name="classification") init_batch = { "inputs": jnp.ones((1, max_seq_length, hidden_dim), jnp.float32) } params = init_layer_variables(rng, classification_projection, init_batch)["params"] expected_keys = {"output_bias", "output_kernel"} self.assertEqual(params.keys(), expected_keys) rng, init_rng = jax.random.split(rng) inputs = jax.random.randint(init_rng, (batch_size, max_seq_length, hidden_dim), minval=0, maxval=10) outputs = classification_projection.apply({"params": params}, inputs=inputs) self.assertEqual(outputs.shape, (batch_size, max_seq_length, NUM_CLASSES))
def __call__( self, input_ids, input_mask, type_ids, labels = None, deterministic = False ): """Applies model for sequence classification. Args: input_ids: <int>[BATCH_SIZE, MAX_SEQ_LENGTH] tokenized inputs. input_mask: <bool>[BATCH_SIZE, MAX_SEQ_LENGTH] mask separating actual inputs from padding. Only used by BERT. type_ids: <int>[BATCH_SIZE, MAX_SEQ_LENGTH] Ids partitioning input into different types. labels: True labels associated with inputs. Generally only required for training. Shape depends on task type: * Classification: <int>[BATCH_SIZE] * Regression: <float>[BATCH_SIZE] deterministic: Whether or not to apply dropout to input. Returns: * If labels supplied (training mode): Model loss and metrics. * If no labels supplied (prediction / evaluation mode): Logits of shape <float>[BATCH_SIZE, n_classes]. """ _, pooled_output = EncoderModel( self.config, name="encoder")( input_ids, input_mask, type_ids, deterministic=deterministic) logits = layers.OutputProjection( n_out=self.n_classes, kernel_init=default_kernel_init, name="classification")( pooled_output) if labels is None: # Code path used during evaluation or prediction; metrics can be computed # from logits by the caller. return logits # Code path used during training. if self.config.dataset_name == "glue/stsb": # Regression task loss = jnp.mean((logits[Ellipsis, 0] - labels)**2) return {"loss": loss, "num_labels": labels.size} else: # Classification task logits = nn.log_softmax(logits) loss = -jnp.mean( jnp.sum(onehot(labels, logits.shape[-1]) * logits, axis=-1)) correct_predictions = jnp.sum(logits.argmax(-1) == labels) return { "loss": loss, "correct_predictions": correct_predictions, "num_labels": labels.size }