def get_dense_span_ends_from_starts(dense_span_starts: tf.Tensor, dense_span_ends: tf.Tensor) -> tf.Tensor: """For every mention start positions finds the corresponding end position.""" seq_len = tf.shape(dense_span_starts)[0] start_pos = tf.cast(tf.where(tf.equal(dense_span_starts, 1)), tf.int32) end_pos = tf.cast(tf.squeeze(tf.where(tf.equal(dense_span_ends, 1)), 1), tf.int32) dense_span_ends_from_starts = tf.zeros(seq_len, dtype=tf.int32) dense_span_ends_from_starts = tf.tensor_scatter_nd_add( dense_span_ends_from_starts, start_pos, end_pos) return dense_span_ends_from_starts
def _do_update(x_update_diff_norm_sq, x_update, hess_matmul_x_update): # pylint: disable=missing-docstring hessian_column_with_l2 = sparse_or_dense_matvecmul( hessian_unregularized_loss_outer, hessian_unregularized_loss_middle * _sparse_or_dense_matmul_onehot( hessian_unregularized_loss_outer, coord), adjoint_a=True) if l2_regularizer is not None: hessian_column_with_l2 += _one_hot_like( hessian_column_with_l2, coord, on_value=2. * l2_regularizer) # Move the batch dimensions of `hessian_column_with_l2` to rightmost in # order to conform to `hess_matmul_x_update`. n = tf.rank(hessian_column_with_l2) perm = tf.roll(tf.range(n), shift=1, axis=0) hessian_column_with_l2 = tf.transpose(a=hessian_column_with_l2, perm=perm) # Update the entire batch at `coord` even if `delta` may be 0 at some # batch coordinates. In those cases, adding `delta` is a no-op. x_update = tf.tensor_scatter_nd_add(x_update, [[coord]], [delta]) with tf.control_dependencies([x_update]): x_update_diff_norm_sq_ = x_update_diff_norm_sq + delta**2 hess_matmul_x_update_ = (hess_matmul_x_update + delta * hessian_column_with_l2) # Hint that loop vars retain the same shape. x_update_diff_norm_sq_.set_shape( x_update_diff_norm_sq_.shape.merge_with( x_update_diff_norm_sq.shape)) hess_matmul_x_update_.set_shape( hess_matmul_x_update_.shape.merge_with( hess_matmul_x_update.shape)) return [ x_update_diff_norm_sq_, x_update, hess_matmul_x_update_ ]
def loop_body(auc_prev, p0_prev, p1_prev, i, j): """Body of the loop to integrate over the ROC/PR curves.""" # We will align intervals from q0 and q1 by moving from end to beginning, # stopping with whatever boundary we encounter first. This boundary is # `b`. We find `a` by continuing the search from `b` to the left. b, i, j = _get_endpoint_b(i, j, q1, q0, batch_shape) # Get q1[i] and q1[i+1] ind = tf.where(tf.greater_equal(i, 0)) q1_i = _get_q_slice(q1, i, ind, batch_shape=batch_shape) ip1 = tf.minimum(i + 1, n_q1 - 1) q1_ip1 = _get_q_slice(q1, ip1, ind, batch_shape=batch_shape) # Get q0[j] and q0[j+1] ind = tf.where(tf.greater_equal(j, 0)) q0_j = _get_q_slice(q0, j, ind, batch_shape=batch_shape) jp1 = tf.minimum(j + 1, n_q0 - 1) q0_jp1 = _get_q_slice(q0, jp1, ind, batch_shape=batch_shape) a = _get_endpoint_a(i, j, q1_i, q0_j, batch_shape) # Calculate the proportion [a, b) represents of [q1[i], q1[i+1]). d1 = _get_interval_proportion(i, n_q1, q1_i, q1_ip1, a, b, batch_shape) # Calculate the proportion [a, b) represents of [q1[i], q1[i+1]). d0 = _get_interval_proportion(j, n_q0, q0_j, q0_jp1, a, b, batch_shape) # Notice that because we assumed within bucket values are distributed # uniformly, we end up with something which is identical to the # trapezoidal rule: definite_integral += (b - a) * (f(a) + f(b)) / 2. if curve == 'ROC': auc = auc_prev + d0 * (p1_prev + d1 / 2.) else: total_scaled_delta = n0 * d0 + n1 * d1 total_scaled_cdf_at_b = n0 * p0_prev + n1 * p1_prev def get_auprc_update(): proportion = (total_scaled_cdf_at_b / (total_scaled_cdf_at_b + total_scaled_delta)) definite_integral = ( (n1 / tf.square(total_scaled_delta)) * (d1 * total_scaled_delta + tf.math.log(proportion) * (d1 * total_scaled_cdf_at_b - p1_prev * total_scaled_delta))) return d1 * definite_integral # Values should be non-negative and we use > 0.0 rather than != 0.0 to # catch possible numerical imprecision. delta_gt_0 = tf.greater(total_scaled_delta, 0.) cdf_gt_0 = tf.greater(total_scaled_cdf_at_b, 0.) d1_gt_0 = tf.greater(d1, 0.) ind = tf.where(delta_gt_0 & cdf_gt_0 & d1_gt_0) auc_update = tf.gather_nd(get_auprc_update(), ind) auc = tf.tensor_scatter_nd_add(auc_prev, ind, auc_update) # TODO(jvdillon): In calculating AUROC and AUPRC, we probably should # resolve ties randomly, making the following eight states possible (where # e = 1 means we expected a positive trial and e = 0 means we expected a # negative trial): # # P(y = 1 | pi(x) > delta) # P(y = 1 | pi(x) = delta, e = 1) 0.5 # P(y = 1 | pi(x) = delta, e = 0) 0.5 # P(y = 1 | pi(x) < delta) # # P(y = 0 | pi(x) > delta) # P(y = 0 | pi(x) = delta, e = 1) 0.5 # P(y = 0 | pi(x) = delta, e = 0) 0.5 # P(y = 0 | pi(x) < delta) # # This makes the math a bit harder and its not clear this adds much, # especially when we're assuming piecewise uniformity. # Accumulate this mass (d1, d0) for the next iteration. p1 = p1_prev + d1 p0 = p0_prev + d0 return auc, p0, p1, i, j
def collater_fn(batch: Dict[Text, tf.Tensor]) -> Dict[Text, tf.Tensor]: """Collater function for mention classification task. See BaseTask.""" new_batch = {} # Sample mentions uniformly across batch mention_mask = tf.reshape(batch['mention_mask'], [n_candidate_mentions]) sample_scores = tf.random.uniform( shape=[n_candidate_mentions]) * tf.cast( mention_mask, tf.float32) mention_target_indices = tf.reshape( batch['mention_target_indices'], [bsz]) # We want to make sure that the target mentions always have a priority # when we sample `max_batch_mentions` out of all available mentions. # Additionally, we want these target mentions to be in the same order as # their samples. In other words, we want the first sampled mention to be # target mention from the first sample, the second sampled mention to be # tagret mention from the second sample, etc. # Positions of target mentions in the flat array mention_target_indices_flat = (tf.cast( tf.range(bsz) * max_mentions_per_sample, mention_target_indices.dtype) + mention_target_indices) # These extra score makes sure that target mentions have a priority and # will be sampled in the correct order. mention_target_extra_score_flat = tf.cast( tf.reverse(tf.range(bsz) + 1, axis=[0]), tf.float32) # The model assumes that there is only ONE target mention per sample. # Moreover,we want to select them according to the order of samples: # target mention from sample 0, target mention from sample 1, ..., etc. sample_scores = tf.tensor_scatter_nd_add( sample_scores, tf.expand_dims(mention_target_indices_flat, 1), mention_target_extra_score_flat) sampled_indices = tf.math.top_k(sample_scores, max_batch_mentions, sorted=True).indices # Double-check target mentions were selected correctly. assert_op = tf.assert_equal( sampled_indices[:bsz], tf.cast(mention_target_indices_flat, sampled_indices.dtype)) with tf.control_dependencies([assert_op]): mention_mask = tf.gather(mention_mask, sampled_indices) dtype = batch['mention_start_positions'].dtype mention_start_positions = tf.gather( tf.reshape(batch['mention_start_positions'], [n_candidate_mentions]), sampled_indices) mention_end_positions = tf.gather( tf.reshape(batch['mention_end_positions'], [n_candidate_mentions]), sampled_indices) mention_batch_positions = tf.gather( tf.repeat(tf.range(bsz, dtype=dtype), max_mentions_per_sample), sampled_indices) new_batch['text_ids'] = batch['text_ids'] new_batch['text_mask'] = batch['text_mask'] new_batch['classifier_target'] = tf.reshape( batch['target'], [bsz, config.max_num_labels_per_sample]) new_batch['classifier_target_mask'] = tf.reshape( batch['target_mask'], [bsz, config.max_num_labels_per_sample]) new_batch['mention_mask'] = mention_mask new_batch['mention_start_positions'] = mention_start_positions new_batch['mention_end_positions'] = mention_end_positions new_batch['mention_batch_positions'] = mention_batch_positions new_batch['mention_target_indices'] = tf.range(bsz, dtype=dtype) if config.get('max_length_with_entity_tokens') is not None: batch_with_entity_tokens = mention_preprocess_utils.add_entity_tokens( text_ids=new_batch['text_ids'], text_mask=new_batch['text_mask'], mention_mask=new_batch['mention_mask'], mention_batch_positions=new_batch[ 'mention_batch_positions'], mention_start_positions=new_batch[ 'mention_start_positions'], mention_end_positions=new_batch['mention_end_positions'], new_length=config.max_length_with_entity_tokens, ) # Update `text_ids`, `text_mask`, `mention_mask`, `mention_*_positions` new_batch.update(batch_with_entity_tokens) # Update `max_length` max_length = config.max_length_with_entity_tokens else: max_length = encoder_config.max_length new_batch['mention_target_batch_positions'] = tf.gather( new_batch['mention_batch_positions'], new_batch['mention_target_indices']) new_batch['mention_target_start_positions'] = tf.gather( new_batch['mention_start_positions'], new_batch['mention_target_indices']) new_batch['mention_target_end_positions'] = tf.gather( new_batch['mention_end_positions'], new_batch['mention_target_indices']) new_batch['mention_target_weights'] = tf.ones(bsz) # Fake IDs -- some encoders (ReadTwice) need them new_batch['mention_target_ids'] = tf.zeros(bsz) new_batch['segment_ids'] = tf.zeros_like(new_batch['text_ids']) position_ids = tf.expand_dims(tf.range(max_length, dtype=dtype), axis=0) new_batch['position_ids'] = tf.tile(position_ids, (bsz, 1)) return new_batch
def add_entity_tokens( text_ids: tf.Tensor, text_mask: tf.Tensor, mention_mask: tf.Tensor, mention_batch_positions: tf.Tensor, mention_start_positions: tf.Tensor, mention_end_positions: tf.Tensor, new_length: int, mlm_target_positions: Optional[tf.Tensor] = None, mlm_target_weights: Optional[tf.Tensor] = None, entity_start_token_id: int = default_values.ENTITY_START_TOKEN, entity_end_token_id: int = default_values.ENTITY_END_TOKEN, ) -> Dict[str, tf.Tensor]: """Adds entity start / end tokens around mentions. Inserts `entity_start_token_id` and `entity_end_token_id` tokens around each mention and update mention_start_positions / mention_end_positions to point to these tokens. New text length will be `new_length` and texts will be truncated if nessesary. If a mention no longer fits into the new text, its mask (`mention_mask`) will be set to 0. The function can also update MLM position and weights (`mlm_target_positions` and `mlm_target_weights`) if these arguments are provided. Similarly to mentions, if MLM position no longer fits into the new text, its mask (`mlm_target_weights`) will be set to 0. Args: text_ids: [seq_length] tensor with token ids. text_mask: [seq_length] tensor with 1s for tokens and 0 for padding. mention_mask: [n_mentions] mask indicating whether a mention is a padding. mention_batch_positions: [n_mentions] sample ID of a mention in the batch. mention_start_positions: [n_mentions] position of a mention first token within a sample. mention_end_positions: [n_mentions] position of a mention last token within a sample. new_length: new length of text after entity tokens will be added. mlm_target_positions: [batch_size, max_mlm_targets] positions of tokens to be used for MLM task. mlm_target_weights: [batch_size, max_mlm_targets] mask indicating whether `mlm_target_positions` is a padding. entity_start_token_id: token to be used as entity start token. entity_end_token_id: token to be used as entity end token. Returns: New text_ids and text_mask, updated mentions positions including mention_start_positions, mention_end_positions and mention_mask. Returns updated mlm_target_positions and mlm_target_weights if they were provided as arguments. """ batch_size = tf.shape(text_ids)[0] old_length = tf.shape(text_ids)[1] new_shape = (batch_size, new_length) mentions_fit_mask = compute_which_mentions_fit_with_entity_tokens( mention_mask, mention_batch_positions, mention_start_positions, mention_end_positions, batch_size, old_length, new_length, ) # Ignore mentions that does not fit into new texts. new_mention_mask = mention_mask * mentions_fit_mask mention_start_positions = mention_start_positions * new_mention_mask mention_end_positions = mention_end_positions * new_mention_mask positions = compute_positions_shift_with_entity_tokens( new_mention_mask, mention_batch_positions, mention_start_positions, mention_end_positions, batch_size, old_length) def get_2d_index(positions: tf.Tensor) -> tf.Tensor: return _get_2d_index(mention_batch_positions, positions) def get_new_positions(old_positions: tf.Tensor) -> tf.Tensor: index_2d = get_2d_index(old_positions) return tf.gather_nd(positions, index_2d) new_mention_start_positions = get_new_positions( mention_start_positions) - 1 new_mention_start_positions = new_mention_start_positions * new_mention_mask new_mention_end_positions = get_new_positions(mention_end_positions) + 1 new_mention_end_positions = new_mention_end_positions * new_mention_mask if mlm_target_positions is not None: if mlm_target_weights is None: raise ValueError('`mlm_target_weights` must be specified if ' '`mlm_target_positions` is provided.') mlm_target_positions = tf.gather(positions, mlm_target_positions, batch_dims=1) mlm_target_positions_within_len = tf.less(mlm_target_positions, new_length) mlm_target_positions_within_len = tf.cast( mlm_target_positions_within_len, mlm_target_weights.dtype) mlm_target_weights = mlm_target_weights * mlm_target_positions_within_len # Zero-out positions for pad MLM targets mlm_target_positions = mlm_target_positions * mlm_target_weights # Cut texts that are longer than `new_length` text_within_new_length = tf.less(positions, new_length) text_ids = text_ids * tf.cast(text_within_new_length, text_ids.dtype) text_mask = text_mask * tf.cast(text_within_new_length, text_mask.dtype) positions = tf.minimum(positions, new_length - 1) # Prepare 2D index for tokens positions in the next text_ids and text_mask. # Note that we use flat 2D index and flat values # (e.g. `tf.reshape(text_ids, [-1])`) since `tf.scatter_nd` does not support # batch dimension. batch_positions = _batched_range(old_length, batch_size, 1, positions.dtype) batch_positions = tf.reshape(batch_positions, [-1]) text_index_2d = _get_2d_index(batch_positions, tf.reshape(positions, [-1])) new_text_ids = tf.scatter_nd(text_index_2d, tf.reshape(text_ids, [-1]), new_shape) new_text_mask = tf.scatter_nd(text_index_2d, tf.reshape(text_mask, [-1]), new_shape) # Insert entity start / end tokens into the new text_ids and text_mask. new_mention_start_positions_2d = get_2d_index(new_mention_start_positions) new_mention_end_positions_2d = get_2d_index(new_mention_end_positions) new_text_ids = tf.tensor_scatter_nd_add( new_text_ids, new_mention_start_positions_2d, new_mention_mask * entity_start_token_id) new_text_ids = tf.tensor_scatter_nd_add( new_text_ids, new_mention_end_positions_2d, new_mention_mask * entity_end_token_id) new_mention_mask = tf.cast(new_mention_mask, dtype=text_mask.dtype) new_text_mask = tf.tensor_scatter_nd_add(new_text_mask, new_mention_start_positions_2d, new_mention_mask) new_text_mask = tf.tensor_scatter_nd_add(new_text_mask, new_mention_end_positions_2d, new_mention_mask) features = { 'text_ids': new_text_ids, 'text_mask': new_text_mask, 'mention_start_positions': new_mention_start_positions, 'mention_end_positions': new_mention_end_positions, 'mention_mask': new_mention_mask, } if mlm_target_positions is not None: features['mlm_target_weights'] = mlm_target_weights features['mlm_target_positions'] = mlm_target_positions return features
def _solve_dirichlet(self): # check the boundary conditions are correct if self.domain.boundary.boundary_condition_type != 'Dirichlet': raise ValueError( "_solve_dirichlet must be used with Dirichlet boundary conditions." ) A, b = self._assemble() # work out if any batching needs to be done if len(self.batch_shape) == 0: # add an extra leading dimension onto A A = A[tf.newaxis, ...] batch_shape = [1] # false batch shape, squeezed out by end else: batch_shape = self.batch_shape # get the interior of A global_stiffness_interior_indices = [ *itertools.product(self.mesh.interior_node_indices, repeat=2) ] Ainterior = tf.map_fn( lambda x: tf.reshape( tf.gather_nd(x, global_stiffness_interior_indices), [ len(self.mesh.interior_node_indices), len(self.mesh.interior_node_indices) ]), A) b_interior = tf.gather_nd(b, [ *zip(self.mesh.interior_node_indices, [0] * len(self.mesh.interior_node_indices)) ]) interior_bound_indices = [ *itertools.product(self.mesh.interior_node_indices, self.mesh.boundary_node_indices) ] Aint_bnd = tf.map_fn( lambda x: tf.reshape(tf.gather_nd(x, interior_bound_indices), [ len(self.mesh.interior_node_indices), len(self.mesh.boundary_node_indices) ]), A) bnd_node_indices = np.array(self.mesh.boundary_node_indices, dtype=np.intp) int_node_indices = np.array(self.mesh.interior_node_indices, dtype=np.intp) # get the value on the boundary g = self.domain.boundary.g # convert the stiffness matrices to operators for batch matmul Ainterior_op = tf.linalg.LinearOperatorFullMatrix(Ainterior) Aint_bnd_op = tf.linalg.LinearOperatorFullMatrix(Aint_bnd) # add the fixed dirichlet conditions to sol # ToDo: Batch boundary values sol = tf.scatter_nd(bnd_node_indices[:, None], g, shape=[self.mesh.n_nodes, 1]) b_ = b_interior[..., tf.newaxis] - Aint_bnd_op.matmul(g) sol_interior = Ainterior_op.solve(b_) # sol_interior has a batched shape [b, n_interior_nodes, 1] return tf.squeeze( tf.map_fn( lambda x: tf.tensor_scatter_nd_add( sol, int_node_indices[:, None], x), sol_interior) )[..., tf. newaxis] # kills pesduo-batch dimensions, but keeps output a vector