jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def update_state(self, y_true: tf.Tensor, y_pred: tf.Tensor):
    """Updates metric state.

    Args:
      y_true: `dict`, dictionary with the following name, and key values.
        - masks: [batch, height, width, num_layers], ground-truth masks. The
          num_layers is 1 by default, while all the operations in this function
          support num_layers > 1.
        - valid_masks: [batch, height, width, num_layers], valid elements in the
          mask.
        - image_info: [batch, 4, 2], a tensor that holds information about
          original and preprocessed images. Each entry is in the format of
          [[original_height, original_width], [input_height, input_width],
          [y_scale, x_scale], [y_offset, x_offset]], where [desired_height,
          desired_width] is the actual scaled image size, and [y_scale, x_scale]
          is the scaling factor, which is the ratio of scaled dimension /
          original dimension.
      y_pred: Tensor [batch, height_p, width_p, num_classes], predicated masks.
    """
    logits, gt_masks, valid_masks = preprocess_inputs(y_true, y_pred,
                                                      self._rescale_predictions)
    valid_masks = tf.to_int32(valid_masks)

    gt_binary_masks = tf.one_hot(
        tf.to_int32(gt_masks[..., 0]),
        depth=self.num_classes,
        on_value=True,
        off_value=False,
    )
    gt_binary_masks &= valid_masks

    predictions_binary_masks = tf.one_hot(
        tf.argmax(logits, axis=-1, output_type=tf.int32),
        depth=self.num_classes,
        on_value=True,
        off_value=False,
    )
    predictions_binary_masks &= valid_masks

    super().update_state(
        y_true=gt_binary_masks, y_pred=predictions_binary_masks
    )
//  [python]
def update_state(self, y_true: tf.Tensor, y_pred: tf.Tensor):
    """Updates metric state.

    Args:
      y_true: `dict`, dictionary with the following name, and key values.
        - masks: [batch, height, width, num_layers], ground-truth masks. The
          num_layers is 1 by default, while all the operations in this function
          support num_layers > 1.
        - valid_masks: [batch, height, width, num_layers], valid elements in the
          mask.
        - image_info: [batch, 4, 2], a tensor that holds information about
          original and preprocessed images. Each entry is in the format of
          [[original_height, original_width], [input_height, input_width],
          [y_scale, x_scale], [y_offset, x_offset]], where [desired_height,
          desired_width] is the actual scaled image size, and [y_scale, x_scale]
          is the scaling factor, which is the ratio of scaled dimension /
          original dimension.
      y_pred: Tensor [batch, height_p, width_p, num_classes], predicated masks.
    """
    logits, gt_masks, valid_masks = preprocess_inputs(y_true, y_pred,
                                                      self._rescale_predictions)
    valid_masks = tf.cast(valid_masks, tf.bool)

    gt_binary_masks = tf.one_hot(
        tf.cast(gt_masks[..., 0], dtype=tf.int32),
        depth=self.num_classes,
        on_value=True,
        off_value=False,
    )
    gt_binary_masks &= valid_masks

    predictions_binary_masks = tf.one_hot(
        tf.argmax(logits, axis=-1, output_type=tf.int32),
        depth=self.num_classes,
        on_value=True,
        off_value=False,
    )
    predictions_binary_masks &= valid_masks

    super().update_state(
        y_true=gt_binary_masks, y_pred=predictions_binary_masks
    )
//  [python]
def CalculateReceptiveBoxes(height, width, rf, stride, padding):
  """Calculate receptive boxes for each feature point.

  Args:
    height: The height of feature map.
    width: The width of feature map.
    rf: The receptive field size.
    stride: The effective stride between two adjacent feature points.
    padding: The effective padding size.

  Returns:
    rf_boxes: [N, 4] receptive boxes tensor. Here N equals to height x width.
    Each box is represented by [ymin, xmin, ymax, xmax].
  """
  x, y = tf.meshgrid(tf.range(width), tf.range(height))
  coordinates = tf.reshape(tf.stack([y, x], axis=2), [-1, 2])
  # [y,x,y,x]
  point_boxes = tf.concat([coordinates, coordinates], 1)
  bias = [-padding, -padding, -padding + rf - 1, -padding + rf - 1]
  rf_boxes = stride * point_boxes + bias
  return rf_boxes
//  [python]
def CalculateReceptiveBoxes(height, width, rf, stride, padding):
  """Calculate receptive boxes for each feature point.

  Args:
    height: The height of feature map.
    width: The width of feature map.
    rf: The receptive field size.
    stride: The effective stride between two adjacent feature points.
    padding: The effective padding size.

  Returns:
    rf_boxes: [N, 4] receptive boxes tensor. Here N equals to height x width.
    Each box is represented by [ymin, xmin, ymax, xmax].
  """
  x, y = tf.meshgrid(tf.range(width), tf.range(height))
  coordinates = tf.reshape(tf.stack([y, x], axis=2), [-1, 2])
  # [y,x,y,x]
  point_boxes = tf.cast(
      tf.concat([coordinates, coordinates], 1), dtype=tf.float32)
  bias = [-padding, -padding, -padding + rf - 1, -padding + rf - 1]
  rf_boxes = stride * point_boxes + bias
  return rf_boxes
//  [python]
def _rescale_and_pad_predictions(
    predictions: tf.Tensor, images_info: tf.Tensor,
    output_size: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor]:
  """Scales back predictions to original image shapes and pads to output size.

  Args:
    predictions: A tensor in shape [batch, height, width, num_classes] which
      stores the model predictions.
    images_info: A tensor in shape [batch, 4, 2] that holds information about
      original and preprocessed images. Each entry is in the format of
      [[original_height, original_width], [input_height, input_width], [y_scale,
      x_scale], [y_offset, x_offset]], where [desired_height, desired_width] is
      the actual scaled image size, and [y_scale, x_scale] is the scaling
      factor, which is the ratio of scaled dimension / original dimension.
    output_size: A list/tuple/tensor stores the size of the padded output in
      [output_height, output_width].

  Returns:
    predictions: A tensor in shape [batch, output_height, output_width,
      num_classes] which stores the rescaled and padded predictions.
    image_shape_masks: A bool tensor in shape [batch, output_height,
      output_width] where the pixels inside the original image shape are true,
      otherwise false.
  """
  # (batch_size, 2)
  image_shape = tf.to_int32(images_info[:, 0, :])
  desired_size = tf.to_float(images_info[:, 1, :])
  image_scale = tf.to_float(images_info[:, 2, :])
  offset = tf.to_int32(images_info[:, 3, :])
  rescale_size = tf.to_int32(tf.math.ceil(desired_size / image_scale))

  # Rescale the predictions, then crop to the original image shape and
  # finally pad zeros to match the mask size.
  predictions = (
      spatial_transform_ops.bilinear_resize_with_crop_and_pad(
          predictions,
          rescale_size,
          crop_offset=offset,
          crop_size=image_shape,
          output_size=output_size))

  # (batch_size, 2)
  y0_x0 = tf.broadcast_to(
      tf.constant([[0, 0]], dtype=image_shape.dtype), tf.shape(image_shape))
  # (batch_size, 4)
  image_shape_bbox = tf.concat([y0_x0, image_shape], axis=1)
  # (batch_size, height, width)
  image_shape_masks = box_ops.bbox2mask(
      bbox=image_shape_bbox,
      image_height=output_size[0],
      image_width=output_size[1],
      dtype=tf.bool)

  return predictions, image_shape_masks
//  [python]
def _rescale_and_pad_predictions(
    predictions: tf.Tensor, images_info: tf.Tensor,
    output_size: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor]:
  """Scales back predictions to original image shapes and pads to output size.

  Args:
    predictions: A tensor in shape [batch, height, width, num_classes] which
      stores the model predictions.
    images_info: A tensor in shape [batch, 4, 2] that holds information about
      original and preprocessed images. Each entry is in the format of
      [[original_height, original_width], [input_height, input_width], [y_scale,
      x_scale], [y_offset, x_offset]], where [desired_height, desired_width] is
      the actual scaled image size, and [y_scale, x_scale] is the scaling
      factor, which is the ratio of scaled dimension / original dimension.
    output_size: A list/tuple/tensor stores the size of the padded output in
      [output_height, output_width].

  Returns:
    predictions: A tensor in shape [batch, output_height, output_width,
      num_classes] which stores the rescaled and padded predictions.
    image_shape_masks: A bool tensor in shape [batch, output_height,
      output_width] where the pixels inside the original image shape are true,
      otherwise false.
  """
  # (batch_size, 2)
  image_shape = tf.cast(images_info[:, 0, :], tf.int32)
  desired_size = tf.cast(images_info[:, 1, :], tf.float32)
  image_scale = tf.cast(images_info[:, 2, :], tf.float32)
  offset = tf.cast(images_info[:, 3, :], tf.int32)
  rescale_size = tf.cast(tf.math.ceil(desired_size / image_scale), tf.int32)

  # Rescale the predictions, then crop to the original image shape and
  # finally pad zeros to match the mask size.
  predictions = (
      spatial_transform_ops.bilinear_resize_with_crop_and_pad(
          predictions,
          rescale_size,
          crop_offset=offset,
          crop_size=image_shape,
          output_size=output_size))

  # (batch_size, 2)
  y0_x0 = tf.broadcast_to(
      tf.constant([[0, 0]], dtype=image_shape.dtype), tf.shape(image_shape))
  # (batch_size, 4)
  image_shape_bbox = tf.concat([y0_x0, image_shape], axis=1)
  # (batch_size, height, width)
  image_shape_masks = box_ops.bbox2mask(
      bbox=image_shape_bbox,
      image_height=output_size[0],
      image_width=output_size[1],
      dtype=tf.bool)

  return predictions, image_shape_masks
//  [python]
def _update_cache(self, key, value, cache, decode_loop_step):
    """Updates cache states and gets full-length key/value tensors."""
    # Combines cached keys and values with new keys and values.
    if decode_loop_step is not None:
      # TPU special case.
      key_seq_dim = cache["key"].shape.as_list()[1]
      indices = tf.reshape(
          tf.one_hot(decode_loop_step, key_seq_dim, dtype=key.dtype),
          [1, key_seq_dim, 1, 1])
      key = cache["key"] + key * indices
      value_seq_dim = cache["value"].shape.as_list()[1]
      indices = tf.reshape(
          tf.one_hot(decode_loop_step, value_seq_dim, dtype=value.dtype),
          [1, value_seq_dim, 1, 1])
      value = cache["value"] + value * indices
    else:
      key = tf.concat([tf.to_float(cache["key"]), key], axis=1)
      value = tf.concat([tf.to_float(cache["value"]), value], axis=1)

    # Update cache
    cache["key"] = key
    cache["value"] = value

    return key, value
//  [python]
def _update_cache(self, key, value, cache, decode_loop_step):
    """Updates cache states and gets full-length key/value tensors."""
    # Combines cached keys and values with new keys and values.
    if decode_loop_step is not None:
      # TPU special case.
      key_seq_dim = cache["key"].shape.as_list()[1]
      indices = tf.reshape(
          tf.one_hot(decode_loop_step, key_seq_dim, dtype=key.dtype),
          [1, key_seq_dim, 1, 1])
      key = cache["key"] + key * indices
      value_seq_dim = cache["value"].shape.as_list()[1]
      indices = tf.reshape(
          tf.one_hot(decode_loop_step, value_seq_dim, dtype=value.dtype),
          [1, value_seq_dim, 1, 1])
      value = cache["value"] + value * indices
    else:
      key = tf.concat([tf.cast(cache["key"], key.dtype), key], axis=1)
      value = tf.concat([tf.cast(cache["value"], value.dtype), value], axis=1)

    # Update cache
    cache["key"] = key
    cache["value"] = value

    return key, value
//  [python]
def __init__(self, features, rnnSeqLengths, NUM_CLASSES, keepProb, weight_decay=1e-5, is_training=True):
        self.features = features
        self.rnnSeqLengths = tf.cast(tf.ceil(tf.divide(rnnSeqLengths, FLAGS.width_stride)), tf.int32)
        self.keepProb = keepProb

        self.NUM_CLASSES = NUM_CLASSES
        self.featureNum_reduceHeight = FLAGS.featureNum_reduceHeight
        self.num_hidden = FLAGS.num_hidden

        self.convLayers(weight_decay, is_training)
        self.lstmLayers(is_training)
//  [python]
def compute_accuracy(v_xs,v_ys,sess):
    #prediction 变为全剧变量
    global prediction
    y_pre = sess.run(prediction,feed_dict={xs:v_xs,keep_prob:1})
    #预测值每行是10列，tf.argmax(数据，axis），相等为1，不想等为0
    correct_prediction = tf.equal(tf.argmax(y_pre,1),tf.argmax(v_ys,1))
    # 计算平均值，即计算准确率
    accuracy = tf.reduce_mean(tf.to_float(correct_prediction))
    # 运行我们的accuracy这一步
    result = sess.run(accuracy,feed_dict={xs:v_xs,ys:v_ys,keep_prob:1})
    return result
//  [python]
def compute_accuracy(v_xs,v_ys,sess):
    #prediction 变为全剧变量
    global prediction
    y_pre = sess.run(prediction,feed_dict={xs:v_xs,keep_prob:1})
    #预测值每行是10列，tf.argmax(数据，axis），相等为1，不想等为0
    correct_prediction = tf.equal(tf.argmax(y_pre,1),tf.argmax(v_ys,1))
    # 计算平均值，即计算准确率
    accuracy = tf.reduce_mean(tf.cast(correct_prediction,tf.float32))
    # 运行我们的accuracy这一步
    result = sess.run(accuracy,feed_dict={xs:v_xs,ys:v_ys,keep_prob:1})
    return result
//  [python]
def image_classifier(features, labels, mode, params):
    model_functions = {
        "linear": linear_model,
        "dnn": dnn_model,
        "dnn_dropout": dnn_dropout_model,
        "cnn": cnn_model}
    model_function = model_functions[params["model"]] 
    ylogits, nclasses = model_function(features["image"], mode, params)

    probabilities = tf.nn.softmax(logits = ylogits)
    class_int = tf.to_int64(x = tf.argmax(input = ylogits, axis = 1))
    class_str = tf.gather(params = LIST_OF_LABELS, indices = tf.to_int32(x = class_int))
  
    if mode == tf.estimator.ModeKeys.TRAIN or mode == tf.estimator.ModeKeys.EVAL:
        # Convert string label to int
        labels_table = tf.contrib.lookup.index_table_from_tensor(mapping = tf.constant(value = LIST_OF_LABELS, dtype = tf.string))
        labels = labels_table.lookup(keys = labels)

        loss = tf.reduce_mean(input_tensor = tf.nn.softmax_cross_entropy_with_logits_v2(logits = ylogits, labels = tf.one_hot(indices = labels, depth = NCLASSES)))
        
        if mode == tf.estimator.ModeKeys.TRAIN:
            # This is needed for batch normalization, but has no effect otherwise
            update_ops = tf.get_collection(key = tf.GraphKeys.UPDATE_OPS)
            with tf.control_dependencies(control_inputs = update_ops):
                train_op = tf.contrib.layers.optimize_loss(
                    loss = loss, 
                    global_step = tf.train.get_global_step(),
                    learning_rate = params["learning_rate"],
                    optimizer = "Adam")
            eval_metric_ops = None
        else:
            train_op = None
            eval_metric_ops =  {"accuracy": tf.metrics.accuracy(labels = labels, predictions = class_int)}
    else:
        loss = None
        train_op = None
        eval_metric_ops = None
 
    return tf.estimator.EstimatorSpec(
        mode = mode,
        predictions = {"probabilities": probabilities, 
                       "classid": class_int, 
                       "class": class_str},
        loss = loss,
        train_op = train_op,
        eval_metric_ops = eval_metric_ops,
        export_outputs = {"classes": tf.estimator.export.PredictOutput(
            {"probabilities": probabilities, 
             "classid": class_int, 
             "class": class_str})}
    )
//  [python]
def image_classifier(features, labels, mode, params):
    model_functions = {
        "linear": linear_model,
        "dnn": dnn_model,
        "dnn_dropout": dnn_dropout_model,
        "cnn": cnn_model}
    model_function = model_functions[params["model"]] 
    ylogits, nclasses = model_function(features["image"], mode, params)

    probabilities = tf.nn.softmax(logits = ylogits)
    class_int = tf.cast(x = tf.argmax(input = ylogits, axis = 1), dtype = tf.uint8)
    class_str = tf.gather(params = LIST_OF_LABELS, indices = tf.cast(x = class_int, dtype = tf.int32))
  
    if mode == tf.estimator.ModeKeys.TRAIN or mode == tf.estimator.ModeKeys.EVAL:
        # Convert string label to int
        labels_table = tf.contrib.lookup.index_table_from_tensor(mapping = tf.constant(value = LIST_OF_LABELS, dtype = tf.string))
        labels = labels_table.lookup(keys = labels)

        loss = tf.reduce_mean(input_tensor = tf.nn.softmax_cross_entropy_with_logits_v2(logits = ylogits, labels = tf.one_hot(indices = labels, depth = NCLASSES)))
        
        if mode == tf.estimator.ModeKeys.TRAIN:
            # This is needed for batch normalization, but has no effect otherwise
            update_ops = tf.get_collection(key = tf.GraphKeys.UPDATE_OPS)
            with tf.control_dependencies(control_inputs = update_ops):
                train_op = tf.contrib.layers.optimize_loss(
                    loss = loss, 
                    global_step = tf.train.get_global_step(),
                    learning_rate = params["learning_rate"],
                    optimizer = "Adam")
            eval_metric_ops = None
        else:
            train_op = None
            eval_metric_ops =  {"accuracy": tf.metrics.accuracy(labels = labels, predictions = class_int)}
    else:
        loss = None
        train_op = None
        eval_metric_ops = None
 
    return tf.estimator.EstimatorSpec(
        mode = mode,
        predictions = {"probabilities": probabilities, 
                       "classid": class_int, 
                       "class": class_str},
        loss = loss,
        train_op = train_op,
        eval_metric_ops = eval_metric_ops,
        export_outputs = {"classes": tf.estimator.export.PredictOutput(
            {"probabilities": probabilities, 
             "classid": class_int, 
             "class": class_str})}
    )
//  [python]
def objective(self, params, data, labels):
    """Computes the softmax cross entropy.

    Args:
      params: List of parameter tensors or variables
      data: Batch of features with samples along the first dimension
      labels: Vector of labels with the same number of samples as the data

    Returns:
      loss: Softmax cross entropy loss averaged over the samples in the batch

    Raises:
      ValueError: If the objective is to be computed over >2 classes, because
        this operation is broken in tensorflow at the moment.
    """
    # Forward pass.
    logits = self.inference(params, data)

    # Compute the loss.
    #l2reg = [tf.reduce_sum(param ** 2) for param in params]
    if int(logits.get_shape()[1]) == 2:
      labels = tf.to_float(labels)
      losses = tf.nn.sigmoid_cross_entropy_with_logits(
          labels=labels, logits=logits[:, 0])
    else:
      labels = tf.squeeze(labels)
      losses = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=logits, labels=labels)

    return tf.reduce_mean(losses)
//  [python]
def objective(self, params, data, labels):
    """Computes the softmax cross entropy.

    Args:
      params: List of parameter tensors or variables
      data: Batch of features with samples along the first dimension
      labels: Vector of labels with the same number of samples as the data

    Returns:
      loss: Softmax cross entropy loss averaged over the samples in the batch

    Raises:
      ValueError: If the objective is to be computed over >2 classes, because
        this operation is broken in tensorflow at the moment.
    """
    # Forward pass.
    logits = self.inference(params, data)

    # Compute the loss.
    #l2reg = [tf.reduce_sum(param ** 2) for param in params]
    if int(logits.get_shape()[1]) == 2:
      labels = tf.cast(labels, tf.float32)
      losses = tf.nn.sigmoid_cross_entropy_with_logits(
          labels=labels, logits=logits[:, 0])
    else:
      labels = tf.squeeze(labels)
      losses = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=logits, labels=labels)

    return tf.reduce_mean(losses)
//  [python]
def prepare_mnist_features_and_labels(x, y):
    x = tf.expand_dims(tf.to_float(x), axis=-1) / 255.0
    y = tf.to_int64(y)
    return x, y
//  [python]
def prepare_mnist_features_and_labels(x, y):
    x = tf.expand_dims(tf.cast(x, tf.float32), axis=-1) / 255.0
    y = tf.cast(y, tf.int64)
    return x, y
//  [python]
def _make_expand_x_fn_for_non_batch_interpolation(y_ref, axis):
  """Make func to expand left/right (of axis) dims of tensors shaped like x."""
  # This expansion is to help x broadcast with `y`, the output.
  # In the non-batch case, the output shape is going to be
  #   y_ref.shape[:axis] + x.shape + y_ref.shape[axis+1:]

  # Recall we made axis non-negative
  y_ref_shape = tf.shape(input=y_ref)
  y_ref_shape_left = y_ref_shape[:axis]
  y_ref_shape_right = y_ref_shape[axis + 1:]

  def expand_ends(x, broadcast=False):
    """Expand x so it can bcast w/ tensors of output shape."""
    # Assume out_shape = A + x.shape + B, and rank(A) = axis.
    # Expand with singletons with same rank as A, B.
    expanded_shape = tf.pad(
        tensor=tf.shape(input=x),
        paddings=[[axis, tf.size(input=y_ref_shape_right)]],
        constant_values=1)
    x_expanded = tf.reshape(x, expanded_shape)

    if broadcast:
      out_shape = tf.concat((
          y_ref_shape_left,
          tf.shape(input=x),
          y_ref_shape_right,
      ),
                            axis=0)
      if x.dtype.is_bool:
        x_expanded = x_expanded | t.zeros(out_shape, tf.bool)
      else:
        x_expanded += t.zeros(out_shape, dtype=x.dtype)
    return x_expanded

  return expand_ends
//  [python]
def _make_expand_x_fn_for_non_batch_interpolation(y_ref, axis):
  """Make func to expand left/right (of axis) dims of tensors shaped like x."""
  # This expansion is to help x broadcast with `y`, the output.
  # In the non-batch case, the output shape is going to be
  #   y_ref.shape[:axis] + x.shape + y_ref.shape[axis+1:]

  # Recall we made axis non-negative
  y_ref_shape = tf.shape(input=y_ref)
  y_ref_shape_left = y_ref_shape[:axis]
  y_ref_shape_right = y_ref_shape[axis + 1:]

  def expand_ends(x, broadcast=False):
    """Expand x so it can bcast w/ tensors of output shape."""
    # Assume out_shape = A + x.shape + B, and rank(A) = axis.
    # Expand with singletons with same rank as A, B.
    expanded_shape = tf.pad(
        tensor=tf.shape(input=x),
        paddings=[[axis, tf.size(input=y_ref_shape_right)]],
        constant_values=1)
    x_expanded = tf.reshape(x, expanded_shape)

    if broadcast:
      out_shape = tf.concat((
          y_ref_shape_left,
          tf.shape(input=x),
          y_ref_shape_right,
      ),
                            axis=0)
      if x.dtype.is_bool:
        x_expanded = x_expanded | tf.cast(tf.zeros(out_shape), tf.bool)
      else:
        x_expanded += tf.zeros(out_shape, dtype=x.dtype)
    return x_expanded

  return expand_ends
//  [python]
def gather_tables(
    X, col_gather, context, bos_id, eos_id, table_position_type, max_row_col_embedding
):
    """
    Prepares all the inputs required for initially feeding to one half of the table model.

    Args:
        X: Input tokens [batch, seq]
        col_gather: the output of get_gather_indices
        context: context (usually start and end of each row and column) to be gathered and passed to the model.
        bos_id: The BOS token id. (scalar)
        eos_id: _The EOS token id. (scalar)
        table_position_type: One of "row_col" or "all" - the type of position embedding used.

    Returns:
        _type_: _description_
    """
    pad_id = max_row_col_embedding - 1
    if table_position_type == "row_col":
        pos_raw = context[:, :, 2:]
        pos_eos = tf.constant([pad_id, pad_id])
    elif table_position_type == "all":
        pos_raw = context
        pos_eos = tf.constant([pad_id, pad_id, pad_id, pad_id])

    return {
        **gather_col_vals(X, col_gather, eos_id, bos_id, pad_val=1234),
        "scatter_vals": col_gather["values"],
        "positions": gather_col_vals(
            t.cast(pos_raw, tf.int32), col_gather, pos_eos, pos_eos, pad_val=1
        )["values"],
    }
//  [python]
def gather_tables(
    X, col_gather, context, bos_id, eos_id, table_position_type, max_row_col_embedding
):
    """
    Prepares all the inputs required for initially feeding to one half of the table model.

    Args:
        X: Input tokens [batch, seq]
        col_gather: the output of get_gather_indices
        context: context (usually start and end of each row and column) to be gathered and passed to the model.
        bos_id: The BOS token id. (scalar)
        eos_id: _The EOS token id. (scalar)
        table_position_type: One of "row_col" or "all" - the type of position embedding used.

    Returns:
        _type_: _description_
    """
    pad_id = max_row_col_embedding - 1
    if table_position_type == "row_col":
        pos_raw = context[:, :, 2:]
        pos_eos = tf.constant([pad_id, pad_id])
    elif table_position_type == "all":
        pos_raw = context
        pos_eos = tf.constant([pad_id, pad_id, pad_id, pad_id])

    return {
        **gather_col_vals(X, col_gather, eos_id, bos_id, pad_val=1234),
        "scatter_vals": col_gather["values"],
        "positions": gather_col_vals(
            tf.cast(pos_raw, tf.int32), col_gather, pos_eos, pos_eos, pad_val=1
        )["values"],
    }
//  [python]
def _make_expand_x_fn_for_batch_interpolation(y_ref, axis):
  """Make func to expand left/right (of axis) dims of tensors shaped like x."""
  # This expansion is to help x broadcast with `y`, the output.
  # In the batch case, the output shape is going to be
  #   Broadcast(y_ref.shape[:axis], x.shape[:-1]) +
  #   x.shape[-1:] +  y_ref.shape[axis+1:]

  # Recall we made axis non-negative
  y_ref_shape = tf.shape(input=y_ref)
  y_ref_shape_left = y_ref_shape[:axis]
  y_ref_shape_right = y_ref_shape[axis + 1:]

  def expand_right_dims(x, broadcast=False):
    """Expand x so it can bcast w/ tensors of output shape."""
    expanded_shape_left = tf.broadcast_dynamic_shape(
        tf.shape(input=x)[:-1],
        tf.ones([tf.size(input=y_ref_shape_left)], dtype=tf.int32))
    expanded_shape = tf.concat(
        (expanded_shape_left, tf.shape(input=x)[-1:],
         tf.ones([tf.size(input=y_ref_shape_right)], dtype=tf.int32)),
        axis=0)
    x_expanded = tf.reshape(x, expanded_shape)
    if broadcast:
      broadcast_shape_left = tf.broadcast_dynamic_shape(
          tf.shape(input=x)[:-1], y_ref_shape_left)
      broadcast_shape = tf.concat(
          (broadcast_shape_left, tf.shape(input=x)[-1:], y_ref_shape_right),
          axis=0)
      if x.dtype.is_bool:
        x_expanded = x_expanded | tf.to_float(tf.zeros(broadcast_shape), tf.bool)
      else:
        x_expanded += tf.zeros(broadcast_shape, dtype=x.dtype)
    return x_expanded

  return expand_right_dims
//  [python]
def _make_expand_x_fn_for_batch_interpolation(y_ref, axis):
  """Make func to expand left/right (of axis) dims of tensors shaped like x."""
  # This expansion is to help x broadcast with `y`, the output.
  # In the batch case, the output shape is going to be
  #   Broadcast(y_ref.shape[:axis], x.shape[:-1]) +
  #   x.shape[-1:] +  y_ref.shape[axis+1:]

  # Recall we made axis non-negative
  y_ref_shape = tf.shape(input=y_ref)
  y_ref_shape_left = y_ref_shape[:axis]
  y_ref_shape_right = y_ref_shape[axis + 1:]

  def expand_right_dims(x, broadcast=False):
    """Expand x so it can bcast w/ tensors of output shape."""
    expanded_shape_left = tf.broadcast_dynamic_shape(
        tf.shape(input=x)[:-1],
        tf.ones([tf.size(input=y_ref_shape_left)], dtype=tf.int32))
    expanded_shape = tf.concat(
        (expanded_shape_left, tf.shape(input=x)[-1:],
         tf.ones([tf.size(input=y_ref_shape_right)], dtype=tf.int32)),
        axis=0)
    x_expanded = tf.reshape(x, expanded_shape)
    if broadcast:
      broadcast_shape_left = tf.broadcast_dynamic_shape(
          tf.shape(input=x)[:-1], y_ref_shape_left)
      broadcast_shape = tf.concat(
          (broadcast_shape_left, tf.shape(input=x)[-1:], y_ref_shape_right),
          axis=0)
      if x.dtype.is_bool:
        x_expanded = x_expanded | tf.cast(tf.zeros(broadcast_shape), tf.bool)
      else:
        x_expanded += tf.zeros(broadcast_shape, dtype=x.dtype)
    return x_expanded

  return expand_right_dims
//  [python]
def call(self,
           lm_output_logits,
           sentence_output_logits,
           lm_label_ids,
           lm_label_weights,
           sentence_labels=None):
    """Implements call() for the layer."""
    lm_label_weights = tf.cast(lm_label_weights, tf.float32)
    lm_output_logits = tf.cast(lm_output_logits, tf.float32)

    lm_prediction_losses = tf_keras.losses.sparse_categorical_crossentropy(
        lm_label_ids, lm_output_logits, from_logits=True)
    lm_numerator_loss = tf.reduce_sum(lm_prediction_losses * lm_label_weights)
    lm_denominator_loss = tf.reduce_sum(lm_label_weights)
    mask_label_loss = tf.math.divide_no_nan(lm_numerator_loss,
                                            lm_denominator_loss)

    if sentence_labels is not None:
      sentence_output_logits = tf.cast(sentence_output_logits, tf.float32)
      sentence_loss = tf_keras.losses.sparse_categorical_crossentropy(
          sentence_labels, sentence_output_logits, from_logits=True)
      sentence_loss = tf.reduce_mean(sentence_loss)
      loss = mask_label_loss + sentence_loss
    else:
      sentence_loss = None
      loss = mask_label_loss

    batch_shape = tf.slice(tf.shape(lm_label_ids), [0], [1])
    # TODO(hongkuny): Avoids the hack and switches add_loss.
    final_loss = tf.fill(batch_shape, loss)

    self._add_metrics(lm_output_logits, lm_label_ids, lm_label_weights,
                      mask_label_loss, sentence_output_logits, sentence_labels,
                      sentence_loss)
    return final_loss
//  [python]
def __init__(self, adjusted_lr: float, lr_warmup_init: float,
               lr_warmup_step: int, total_steps: int):
    """Build a CosineLrSchedule.

    Args:
      adjusted_lr: `float`, The initial learning rate.
      lr_warmup_init: `float`, The warm up learning rate.
      lr_warmup_step: `int`, The warm up step.
      total_steps: `int`, Total train steps.
    """
    super().__init__()
    logging.info('LR schedule method: cosine')
    self.adjusted_lr = adjusted_lr
    self.lr_warmup_init = lr_warmup_init
    self.lr_warmup_step = lr_warmup_step
    self.decay_steps = t.cast(total_steps - lr_warmup_step, t.float32)
//  [python]
def __init__(self, adjusted_lr: float, lr_warmup_init: float,
               lr_warmup_step: int, total_steps: int):
    """Build a CosineLrSchedule.

    Args:
      adjusted_lr: `float`, The initial learning rate.
      lr_warmup_init: `float`, The warm up learning rate.
      lr_warmup_step: `int`, The warm up step.
      total_steps: `int`, Total train steps.
    """
    super().__init__()
    logging.info('LR schedule method: cosine')
    self.adjusted_lr = adjusted_lr
    self.lr_warmup_init = lr_warmup_init
    self.lr_warmup_step = lr_warmup_step
    self.decay_steps = tf.cast(total_steps - lr_warmup_step, tf.float32)
//  [python]
def call(self, y_true, box_outputs):
    num_positives, box_targets = y_true
    normalizer = num_positives * 4.0
    mask = t.cast(box_targets != 0.0, box_outputs.dtype)
    box_targets = tf.expand_dims(box_targets, axis=-1)
    box_outputs = tf.expand_dims(box_outputs, axis=-1)
    # TODO(fsx950223): remove cast when huber loss dtype is fixed.
    box_loss = t.cast(self.huber(box_targets, box_outputs),
                       box_outputs.dtype) * mask
    box_loss = tf.reduce_sum(box_loss) / normalizer
    return box_loss
//  [python]
def call(self, y_true, box_outputs):
    num_positives, box_targets = y_true
    normalizer = num_positives * 4.0
    mask = tf.cast(box_targets != 0.0, box_outputs.dtype)
    box_targets = tf.expand_dims(box_targets, axis=-1)
    box_outputs = tf.expand_dims(box_outputs, axis=-1)
    # TODO(fsx950223): remove cast when huber loss dtype is fixed.
    box_loss = tf.cast(self.huber(box_targets, box_outputs),
                       box_outputs.dtype) * mask
    box_loss = tf.reduce_sum(box_loss) / normalizer
    return box_loss
//  [python]
def call(self, y_true, box_outputs):
    anchor_boxes = tf.tile(
        self.input_anchors.boxes,
        [box_outputs.shape[0] // self.input_anchors.boxes.shape[0], 1])
    num_positives, box_targets = y_true
    normalizer = num_positives * 4.0
    mask = tf.to_float(box_targets != 0.0)
    box_outputs = anchors.decode_box_outputs(box_outputs, anchor_boxes) * mask
    box_targets = anchors.decode_box_outputs(box_targets, anchor_boxes) * mask
    box_iou_loss = iou_utils.iou_loss(box_outputs, box_targets,
                                      self.iou_loss_type)
    box_iou_loss = tf.reduce_sum(box_iou_loss) / normalizer
    return box_iou_loss
//  [python]
def call(self, y_true, box_outputs):
    anchor_boxes = tf.tile(
        self.input_anchors.boxes,
        [box_outputs.shape[0] // self.input_anchors.boxes.shape[0], 1])
    num_positives, box_targets = y_true
    normalizer = num_positives * 4.0
    mask = tf.cast(box_targets != 0.0, box_outputs.dtype)
    box_outputs = anchors.decode_box_outputs(box_outputs, anchor_boxes) * mask
    box_targets = anchors.decode_box_outputs(box_targets, anchor_boxes) * mask
    box_iou_loss = iou_utils.iou_loss(box_outputs, box_targets,
                                      self.iou_loss_type)
    box_iou_loss = tf.reduce_sum(box_iou_loss) / normalizer
    return box_iou_loss
//  [python]
def chunk_ragged_tensor(
    inputs: tf.RaggedTensor,
    other_end: tf.Tensor,
    include_n_rows: int = 2,
    base_model_max_length: int = 512 - 2,
    # -2 to account for BOS and EOS tokens that will be added.
):
    """
    Chunks up inputs, keeping the first include_n_rows in every chunk.
    """
    batch_size = inputs.bounding_shape()[0]

    # Figure out which tokens are in the first include_n_rows rows of the column and split these off.

    is_first_n_rows_orig = tf.gather_nd(
        tf.math.less(other_end, include_n_rows), inputs, batch_dims=0
    )
    less_than_max_length_tokens = tf.math.less(
        tf.math.reduce_sum(tf.cast(is_first_n_rows_orig, tf.int64), axis=1),
        base_model_max_length,
    )
    # If first row tokens is more than base_model_max_length - do not use any first row context.
    is_first_n_rows = tf.logical_and(
        tf.expand_dims(less_than_max_length_tokens, 1), is_first_n_rows_orig
    )
    num_first_n_row_tokens = tf.math.reduce_sum(
        tf.cast(is_first_n_rows, tf.int64), axis=1
    )

    num_cols_per_col = tf.to_int64(
        tf.maximum(
            tf.math.ceil(
                (inputs.row_lengths() - num_first_n_row_tokens)
                / (base_model_max_length - num_first_n_row_tokens)
            ),
            1,  # Max(1 .) is to catch the case where all tokens are in the first n rows.
        ),
    )
    first_n_rows = tf.ragged.boolean_mask(inputs, is_first_n_rows)
    other_rows = tf.ragged.boolean_mask(inputs, tf.math.logical_not(is_first_n_rows))
    other_rows_length_goals = base_model_max_length - num_first_n_row_tokens  # bs

    # Reshape the rest of the columns to the per-column length_goal
    new_row_breaks = tf.to_int64(
        tf.minimum(
            tf.ragged.range(
                (num_cols_per_col + 1) * other_rows_length_goals,
                deltas=other_rows_length_goals,
                dtype=tf.int64,
            ),
            tf.expand_dims(other_rows.row_lengths(), 1),
        ),
    )
    new_row_lengths = new_row_breaks[:, 1:] - new_row_breaks[:, :-1]
    other_rows_reshaped = tf.RaggedTensor.from_row_lengths(
        values=other_rows.flat_values, row_lengths=new_row_lengths.flat_values
    )

    # Tile the first N rows so that they are duplicated for each chunk of the column

    # a really gross way to get range broadcasted to num_cols_per_col - not sure of a better way to get this.
    indexes = tf.ragged.range(num_cols_per_col) * 0 + tf.expand_dims(
        tf.range(batch_size), 1
    )
    first_n_rows_duped = tf.gather(first_n_rows, indexes).merge_dims(0, 1)

    result = tf.concat([first_n_rows_duped, other_rows_reshaped], 1)
    return result
//  [python]
def chunk_ragged_tensor(
    inputs: tf.RaggedTensor,
    other_end: tf.Tensor,
    include_n_rows: int = 2,
    base_model_max_length: int = 512 - 2,
    # -2 to account for BOS and EOS tokens that will be added.
):
    """
    Chunks up inputs, keeping the first include_n_rows in every chunk.
    """
    batch_size = inputs.bounding_shape()[0]

    # Figure out which tokens are in the first include_n_rows rows of the column and split these off.

    is_first_n_rows_orig = tf.gather_nd(
        tf.math.less(other_end, include_n_rows), inputs, batch_dims=0
    )
    less_than_max_length_tokens = tf.math.less(
        tf.math.reduce_sum(tf.cast(is_first_n_rows_orig, tf.int64), axis=1),
        base_model_max_length,
    )
    # If first row tokens is more than base_model_max_length - do not use any first row context.
    is_first_n_rows = tf.logical_and(
        tf.expand_dims(less_than_max_length_tokens, 1), is_first_n_rows_orig
    )
    num_first_n_row_tokens = tf.math.reduce_sum(
        tf.cast(is_first_n_rows, tf.int64), axis=1
    )

    num_cols_per_col = tf.cast(
        tf.maximum(
            tf.math.ceil(
                (inputs.row_lengths() - num_first_n_row_tokens)
                / (base_model_max_length - num_first_n_row_tokens)
            ),
            1,  # Max(1 .) is to catch the case where all tokens are in the first n rows.
        ),
        tf.int64,
    )
    first_n_rows = tf.ragged.boolean_mask(inputs, is_first_n_rows)
    other_rows = tf.ragged.boolean_mask(inputs, tf.math.logical_not(is_first_n_rows))
    other_rows_length_goals = base_model_max_length - num_first_n_row_tokens  # bs

    # Reshape the rest of the columns to the per-column length_goal
    new_row_breaks = tf.cast(
        tf.minimum(
            tf.ragged.range(
                (num_cols_per_col + 1) * other_rows_length_goals,
                deltas=other_rows_length_goals,
                dtype=tf.int64,
            ),
            tf.expand_dims(other_rows.row_lengths(), 1),
        ),
        tf.int64,
    )
    new_row_lengths = new_row_breaks[:, 1:] - new_row_breaks[:, :-1]
    other_rows_reshaped = tf.RaggedTensor.from_row_lengths(
        values=other_rows.flat_values, row_lengths=new_row_lengths.flat_values
    )

    # Tile the first N rows so that they are duplicated for each chunk of the column

    # a really gross way to get range broadcasted to num_cols_per_col - not sure of a better way to get this.
    indexes = tf.ragged.range(num_cols_per_col) * 0 + tf.expand_dims(
        tf.range(batch_size), 1
    )
    first_n_rows_duped = tf.gather(first_n_rows, indexes).merge_dims(0, 1)

    result = tf.concat([first_n_rows_duped, other_rows_reshaped], 1)
    return result
//  [python]
def _coco91_to_80(classif, box, areas, iscrowds):
  """Function used to reduce COCO 91 to COCO 80 (2017 to 2014 format)."""
  # Vector where index i coralates to the class at index[i].
  class_ids = [
      1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22,
      23, 24, 25, 27, 28, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43,
      44, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62,
      63, 64, 65, 67, 70, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 84, 85,
      86, 87, 88, 89, 90
  ]
  new_classes = tf.expand_dims(tf.convert_to_tensor(class_ids), axis=0)

  # Resahpe the classes to in order to build a class mask.
  classes = tf.expand_dims(classif, axis=-1)

  # One hot the classificiations to match the 80 class format.
  ind = classes == new_classes

  # Select the max values.
  selected_class = tf.reshape(
      tf.math.argmax(tf.cast(ind, tf.float32), axis=-1), [-1])
  ind = tf.where(tf.reduce_any(ind, axis=-1))

  # Gather the valuable instances.
  classif = tf.gather_nd(selected_class, ind)
  box = tf.gather_nd(box, ind)
  areas = tf.gather_nd(areas, ind)
  iscrowds = tf.gather_nd(iscrowds, ind)

  # Restate the number of viable detections, ideally it should be the same.
  num_detections = tf.shape(classif)[0]
  return classif, box, areas, iscrowds, num_detections
//  [python]
def _coco91_to_80(classif, box, areas, iscrowds):
  """Function used to reduce COCO 91 to COCO 80 (2017 to 2014 format)."""
  # Vector where index i coralates to the class at index[i].
  class_ids = [
      1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22,
      23, 24, 25, 27, 28, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43,
      44, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62,
      63, 64, 65, 67, 70, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 84, 85,
      86, 87, 88, 89, 90
  ]
  new_classes = tf.expand_dims(tf.convert_to_tensor(class_ids), axis=0)

  # Resahpe the classes to in order to build a class mask.
  classes = tf.expand_dims(classif, axis=-1)

  # One hot the classificiations to match the 80 class format.
  ind = classes == tf.cast(new_classes, classes.dtype)

  # Select the max values.
  selected_class = tf.reshape(
      tf.math.argmax(tf.cast(ind, tf.float32), axis=-1), [-1])
  ind = tf.where(tf.reduce_any(ind, axis=-1))

  # Gather the valuable instances.
  classif = tf.gather_nd(selected_class, ind)
  box = tf.gather_nd(box, ind)
  areas = tf.gather_nd(areas, ind)
  iscrowds = tf.gather_nd(iscrowds, ind)

  # Restate the number of viable detections, ideally it should be the same.
  num_detections = tf.shape(classif)[0]
  return classif, box, areas, iscrowds, num_detections
//  [python]
def accuracy(y_true, logits):
    return tf.reduce_mean(tf.to_float(
            tf.equal(tf.argmax(y_true, 1), tf.argmax(logits, 1))))
//  [python]
def accuracy(y_true, logits):
    return tf.reduce_mean(tf.cast(
            tf.equal(tf.argmax(y_true, 1), tf.argmax(logits, 1)),
            tf.float32))
//  [python]
def _random_mask(self) -> tf.Tensor:
        """Creates a random matrix with `num_ones` 1s and 0s otherwise.

        Returns:
            A random mask matrix
        """
        mask = tf.random.uniform(tf.shape(self.kernel), 0, 1)
        mask = tf.to_int32(tf.math.less(mask, self.density))
        return mask
//  [python]
def _random_mask(self) -> tf.Tensor:
        """Creates a random matrix with `num_ones` 1s and 0s otherwise.

        Returns:
            A random mask matrix
        """
        mask = tf.random.uniform(tf.shape(self.kernel), 0, 1)
        mask = tf.cast(tf.math.less(mask, self.density), self.kernel.dtype)
        return mask
//  [python]
def __call__(self, step):
    linear_warmup = (
        self.lr_warmup_init +
        (t.cast(step, dtype=t.float32) / self.lr_warmup_step *
         (self.adjusted_lr - self.lr_warmup_init)))
    learning_rate = t.where(step < self.lr_warmup_step, linear_warmup,
                             self.adjusted_lr)
    lr_schedule = [[1.0, self.lr_warmup_step], [0.1, self.first_lr_drop_step],
                   [0.01, self.second_lr_drop_step]]
    for mult, start_global_step in lr_schedule:
      learning_rate = t.where(step < start_global_step, learning_rate,
                               self.adjusted_lr * mult)
    return learning_rate
//  [python]
def __call__(self, step):
    linear_warmup = (
        self.lr_warmup_init +
        (tf.cast(step, dtype=tf.float32) / self.lr_warmup_step *
         (self.adjusted_lr - self.lr_warmup_init)))
    learning_rate = tf.where(step < self.lr_warmup_step, linear_warmup,
                             self.adjusted_lr)
    lr_schedule = [[1.0, self.lr_warmup_step], [0.1, self.first_lr_drop_step],
                   [0.01, self.second_lr_drop_step]]
    for mult, start_global_step in lr_schedule:
      learning_rate = tf.where(step < start_global_step, learning_rate,
                               self.adjusted_lr * mult)
    return learning_rate
//  [python]
def _get_bad_mask(
        self, labels: tf.Tensor, target_labels: tf.Tensor, idxs: tf.Tensor
    ) -> tf.Tensor:
        """Calculate bad mask for given indices.

        Checks that input features are different for positive negative samples.
        """
        pos_labels = tf.expand_dims(target_labels, axis=-2)
        neg_labels = layers_utils.get_candidate_values(labels, idxs)

        return tf.to_float(
            tf.reduce_all(tf.equal(neg_labels, pos_labels), axis=-1)
        )
//  [python]
def _get_bad_mask(
        self, labels: tf.Tensor, target_labels: tf.Tensor, idxs: tf.Tensor
    ) -> tf.Tensor:
        """Calculate bad mask for given indices.

        Checks that input features are different for positive negative samples.
        """
        pos_labels = tf.expand_dims(target_labels, axis=-2)
        neg_labels = layers_utils.get_candidate_values(labels, idxs)

        return tf.cast(
            tf.reduce_all(tf.equal(neg_labels, pos_labels), axis=-1), pos_labels.dtype
        )
//  [python]
def call(
        self, logits: tf.Tensor, sequence_lengths: tf.Tensor
    ) -> Tuple[tf.Tensor, tf.Tensor]:
        """Decodes the highest scoring sequence of tags.

        Arguments:
            logits: A [batch_size, max_seq_len, num_tags] tensor of
                unary potentials.
            sequence_lengths: A [batch_size] vector of true sequence lengths.

        Returns:
            A [batch_size, max_seq_len] matrix, with dtype `tf.int32`.
            Contains the highest scoring tag indices.
            A [batch_size, max_seq_len] matrix, with dtype `tf.float32`.
            Contains the confidence values of the highest scoring tag indices.
        """
        predicted_ids, scores, _ = rasa.utils.tensorflow.crf.crf_decode(
            logits, self.transition_params, sequence_lengths
        )
        # set prediction index for padding to `0`
        mask = tf.sequence_mask(
            sequence_lengths,
            maxlen=tf.shape(predicted_ids)[1],
            dtype=predicted_ids.dtype,
        )

        confidence_values = scores * tf.to_float(mask)
        predicted_ids = predicted_ids * mask

        return predicted_ids, confidence_values
//  [python]
def call(
        self, logits: tf.Tensor, sequence_lengths: tf.Tensor
    ) -> Tuple[tf.Tensor, tf.Tensor]:
        """Decodes the highest scoring sequence of tags.

        Arguments:
            logits: A [batch_size, max_seq_len, num_tags] tensor of
                unary potentials.
            sequence_lengths: A [batch_size] vector of true sequence lengths.

        Returns:
            A [batch_size, max_seq_len] matrix, with dtype `tf.int32`.
            Contains the highest scoring tag indices.
            A [batch_size, max_seq_len] matrix, with dtype `tf.float32`.
            Contains the confidence values of the highest scoring tag indices.
        """
        predicted_ids, scores, _ = rasa.utils.tensorflow.crf.crf_decode(
            logits, self.transition_params, sequence_lengths
        )
        # set prediction index for padding to `0`
        mask = tf.sequence_mask(
            sequence_lengths,
            maxlen=tf.shape(predicted_ids)[1],
            dtype=predicted_ids.dtype,
        )

        confidence_values = scores * tf.cast(mask, tf.float32)
        predicted_ids = predicted_ids * mask

        return predicted_ids, confidence_values
//  [python]
def f1_score(
        self, tag_ids: tf.Tensor, pred_ids: tf.Tensor, mask: tf.Tensor
    ) -> tf.Tensor:
        """Calculates f1 score for train predictions."""
        mask_bool = tf.to_float(mask[:, :, 0]) > 0.5

        # pick only non padding values and flatten sequences
        tag_ids_flat = tf.boolean_mask(tag_ids, mask_bool)
        pred_ids_flat = tf.boolean_mask(pred_ids, mask_bool)

        # set `0` prediction to not a prediction
        num_tags = self.num_tags - 1

        tag_ids_flat_one_hot = tf.one_hot(tag_ids_flat - 1, num_tags)
        pred_ids_flat_one_hot = tf.one_hot(pred_ids_flat - 1, num_tags)

        return self.f1_score_metric(tag_ids_flat_one_hot, pred_ids_flat_one_hot)