jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _prepare_image_and_label(self, data):
    """Prepare normalized image and label."""
    height = data['image/height']
    width = data['image/width']

    label = tf.io.decode_image(
        data['image/segmentation/class/encoded'], channels=1
    )
    label = tf.reshape(label, (1, height, width))
    label = tf.to_float(label)

    image = tf.io.decode_image(
        data[self._image_feature.feature_name],
        channels=self._image_feature.num_channels,
        dtype=tf.uint8,
    )
    image = tf.reshape(image, (height, width, self._image_feature.num_channels))
    # Normalizes the image feature with mean and std values, which are divided
    # by 255 because an uint8 image are re-scaled automatically. Images other
    # than uint8 type will be wrongly normalized.
    image = preprocess_ops.normalize_image(
        image,
        [mean / 255.0 for mean in self._image_feature.mean],
        [stddev / 255.0 for stddev in self._image_feature.stddev],
    )

    if self._additional_dense_features:
      input_list = [image]
      for feature_cfg in self._additional_dense_features:
        feature = tf.io.decode_image(
            data[feature_cfg.feature_name],
            channels=feature_cfg.num_channels,
            dtype=tf.uint8,
        )
        feature = tf.reshape(feature, (height, width, feature_cfg.num_channels))
        feature = preprocess_ops.normalize_image(
            feature,
            [mean / 255.0 for mean in feature_cfg.mean],
            [stddev / 255.0 for stddev in feature_cfg.stddev],
        )
        input_list.append(feature)
      concat_input = tf.concat(input_list, axis=2)
    else:
      concat_input = image

    if not self._preserve_aspect_ratio:
      label = tf.reshape(label, [data['image/height'], data['image/width'], 1])
      concat_input = tf.image.resize(
          concat_input, self._output_size, method='bilinear'
      )
      label = tf.image.resize(label, self._output_size, method='nearest')
      label = tf.reshape(label[:, :, -1], [1] + self._output_size)

    return concat_input, label
//  [python]
def _prepare_image_and_label(self, data):
    """Prepare normalized image and label."""
    height = data['image/height']
    width = data['image/width']

    label = tf.io.decode_image(
        data['image/segmentation/class/encoded'], channels=1
    )
    label = tf.reshape(label, (1, height, width))
    label = tf.cast(label, tf.float32)

    image = tf.io.decode_image(
        data[self._image_feature.feature_name],
        channels=self._image_feature.num_channels,
        dtype=tf.uint8,
    )
    image = tf.reshape(image, (height, width, self._image_feature.num_channels))
    # Normalizes the image feature with mean and std values, which are divided
    # by 255 because an uint8 image are re-scaled automatically. Images other
    # than uint8 type will be wrongly normalized.
    image = preprocess_ops.normalize_image(
        image,
        [mean / 255.0 for mean in self._image_feature.mean],
        [stddev / 255.0 for stddev in self._image_feature.stddev],
    )

    if self._additional_dense_features:
      input_list = [image]
      for feature_cfg in self._additional_dense_features:
        feature = tf.io.decode_image(
            data[feature_cfg.feature_name],
            channels=feature_cfg.num_channels,
            dtype=tf.uint8,
        )
        feature = tf.reshape(feature, (height, width, feature_cfg.num_channels))
        feature = preprocess_ops.normalize_image(
            feature,
            [mean / 255.0 for mean in feature_cfg.mean],
            [stddev / 255.0 for stddev in feature_cfg.stddev],
        )
        input_list.append(feature)
      concat_input = tf.concat(input_list, axis=2)
    else:
      concat_input = image

    if not self._preserve_aspect_ratio:
      label = tf.reshape(label, [data['image/height'], data['image/width'], 1])
      concat_input = tf.image.resize(
          concat_input, self._output_size, method='bilinear'
      )
      label = tf.image.resize(label, self._output_size, method='nearest')
      label = tf.reshape(label[:, :, -1], [1] + self._output_size)

    return concat_input, label
//  [python]
def labels_to_text(self, labels):
    """Convert labels to text strings.
    Args:
      labels: int32 tensor with shape [batch_size, max_label_length]
    Returns:
      text: string tensor with shape [batch_size]
    """
    if labels.dtype == tf.int32 or labels.dtype == tf.int64:
      labels = tf.to_int64(labels)
    else:
      raise ValueError('Wrong dtype of labels: {}'.format(labels.dtype))
    chars = self._label_to_char_table.lookup(labels)
    text = tf.reduce_join(chars, axis=1)
    return text
//  [python]
def labels_to_text(self, labels):
    """Convert labels to text strings.
    Args:
      labels: int32 tensor with shape [batch_size, max_label_length]
    Returns:
      text: string tensor with shape [batch_size]
    """
    if labels.dtype == tf.int32 or labels.dtype == tf.int64:
      labels = tf.cast(labels, tf.int64)
    else:
      raise ValueError('Wrong dtype of labels: {}'.format(labels.dtype))
    chars = self._label_to_char_table.lookup(labels)
    text = tf.reduce_join(chars, axis=1)
    return text
//  [python]
def to_absolute_coordinates(keypoints, height, width,
                            check_range=True, scope=None):
  """Converts normalized keypoint coordinates to absolute pixel coordinates.

  This function raises an assertion failed error when the maximum keypoint
  coordinate value is larger than 1.01 (in which case coordinates are already
  absolute).

  Args:
    keypoints: A tensor of shape [num_instances, num_keypoints, 2]
    height: Maximum value for y coordinate of absolute keypoint coordinates.
    width: Maximum value for x coordinate of absolute keypoint coordinates.
    check_range: If True, checks if the coordinates are normalized or not.
    scope: name scope.

  Returns:
    tensor of shape [num_instances, num_keypoints, 2] with absolute coordinates
    in terms of the image size.

  """
  with tf.name_scope(scope, 'ToAbsoluteCoordinates'):
    height = tf.cast(height, tf.float32)
    width = tf.cast(width, tf.float32)

    # Ensure range of input keypoints is correct.
    if check_range:
      max_val = tf.reduce_max(keypoints)
      max_assert = tf.Assert(tf.greater_equal(1.01, max_val),
                             ['maximum keypoint coordinate value is larger '
                              'than 1.01: ', max_val])
      with tf.control_dependencies([max_assert]):
        width = tf.identity(width)

    return scale(keypoints, height, width)
//  [python]
def to_normalized_coordinates(keypoints, height, width,
                              check_range=True, scope=None):
  """Converts absolute keypoint coordinates to normalized coordinates in [0, 1].

  Usually one uses the dynamic shape of the image or conv-layer tensor:
    keypoints = keypoint_ops.to_normalized_coordinates(keypoints,
                                                       tf.shape(images)[1],
                                                       tf.shape(images)[2]),

  This function raises an assertion failed error at graph execution time when
  the maximum coordinate is smaller than 1.01 (which means that coordinates are
  already normalized). The value 1.01 is to deal with small rounding errors.

  Args:
    keypoints: A tensor of shape [num_instances, num_keypoints, 2].
    height: Maximum value for y coordinate of absolute keypoint coordinates.
    width: Maximum value for x coordinate of absolute keypoint coordinates.
    check_range: If True, checks if the coordinates are normalized.
    scope: name scope.

  Returns:
    tensor of shape [num_instances, num_keypoints, 2] with normalized
    coordinates in [0, 1].
  """
  with tf.name_scope(scope, 'ToNormalizedCoordinates'):
    height = tf.to_float(height)
    width = tf.to_float(width)

    if check_range:
      max_val = tf.reduce_max(keypoints)
      max_assert = tf.Assert(tf.greater(max_val, 1.01),
                             ['max value is lower than 1.01: ', max_val])
      with tf.control_dependencies([max_assert]):
        width = tf.identity(width)

    return scale(keypoints, 1.0 / height, 1.0 / width)
//  [python]
def to_normalized_coordinates(keypoints, height, width,
                              check_range=True, scope=None):
  """Converts absolute keypoint coordinates to normalized coordinates in [0, 1].

  Usually one uses the dynamic shape of the image or conv-layer tensor:
    keypoints = keypoint_ops.to_normalized_coordinates(keypoints,
                                                       tf.shape(images)[1],
                                                       tf.shape(images)[2]),

  This function raises an assertion failed error at graph execution time when
  the maximum coordinate is smaller than 1.01 (which means that coordinates are
  already normalized). The value 1.01 is to deal with small rounding errors.

  Args:
    keypoints: A tensor of shape [num_instances, num_keypoints, 2].
    height: Maximum value for y coordinate of absolute keypoint coordinates.
    width: Maximum value for x coordinate of absolute keypoint coordinates.
    check_range: If True, checks if the coordinates are normalized.
    scope: name scope.

  Returns:
    tensor of shape [num_instances, num_keypoints, 2] with normalized
    coordinates in [0, 1].
  """
  with tf.name_scope(scope, 'ToNormalizedCoordinates'):
    height = tf.cast(height, tf.float32)
    width = tf.cast(width, tf.float32)

    if check_range:
      max_val = tf.reduce_max(keypoints)
      max_assert = tf.Assert(tf.greater(max_val, 1.01),
                             ['max value is lower than 1.01: ', max_val])
      with tf.control_dependencies([max_assert]):
        width = tf.identity(width)

    return scale(keypoints, 1.0 / height, 1.0 / width)
//  [python]
def create_fake_data_input_fn(batch_size, features_shape, num_classes):
  """Creates a dummy input function with the given feature and label shapes.

  Args:
    batch_size: integer.
    features_shape: list[int]. Feature shape for an individual example.
    num_classes: integer. Number of labels.

  Returns:
    An input function that is usable in the executor.
  """

  def _dataset_fn(input_context=None):
    """An input function for generating fake data."""
    local_batch_size = input_context.get_per_replica_batch_size(batch_size)
    features = np.random.rand(64, *features_shape)
    labels = np.random.randint(2, size=[64, num_classes])
    # Convert the inputs to a Dataset.
    dataset = tf.data.Dataset.from_tensor_slices((features, labels))
    dataset = dataset.shard(input_context.num_input_pipelines,
                            input_context.input_pipeline_id)

    def _assign_dtype(features, labels):
      features = tf.to_float(features)
      labels = tf.to_float(labels)
      return features, labels

    # Shuffle, repeat, and batch the examples.
    dataset = dataset.map(_assign_dtype)
    dataset = dataset.shuffle(64).repeat()
    dataset = dataset.batch(local_batch_size, drop_remainder=True)
    dataset = dataset.prefetch(buffer_size=64)
    return dataset

  return _dataset_fn
//  [python]
def create_fake_data_input_fn(batch_size, features_shape, num_classes):
  """Creates a dummy input function with the given feature and label shapes.

  Args:
    batch_size: integer.
    features_shape: list[int]. Feature shape for an individual example.
    num_classes: integer. Number of labels.

  Returns:
    An input function that is usable in the executor.
  """

  def _dataset_fn(input_context=None):
    """An input function for generating fake data."""
    local_batch_size = input_context.get_per_replica_batch_size(batch_size)
    features = np.random.rand(64, *features_shape)
    labels = np.random.randint(2, size=[64, num_classes])
    # Convert the inputs to a Dataset.
    dataset = tf.data.Dataset.from_tensor_slices((features, labels))
    dataset = dataset.shard(input_context.num_input_pipelines,
                            input_context.input_pipeline_id)

    def _assign_dtype(features, labels):
      features = tf.cast(features, tf.float32)
      labels = tf.cast(labels, tf.float32)
      return features, labels

    # Shuffle, repeat, and batch the examples.
    dataset = dataset.map(_assign_dtype)
    dataset = dataset.shuffle(64).repeat()
    dataset = dataset.batch(local_batch_size, drop_remainder=True)
    dataset = dataset.prefetch(buffer_size=64)
    return dataset

  return _dataset_fn
//  [python]
def _resize_original_image(image, image_shape):
  image = tf.expand_dims(image, 0)
  image = tf.image.resize(
      image, image_shape, method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
  return tf.squeeze(image, 0).numpy()
//  [python]
def _resize_original_image(image, image_shape):
  image = tf.expand_dims(image, 0)
  image = tf.image.resize(
      image, image_shape, method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
  return tf.cast(tf.squeeze(image, 0), tf.uint8)
//  [python]
def get_text(self, ids):
    """Returns a string corresponding to a sequence of character ids.

        Args:
          ids: a tensor with shape [batch_size, max_sequence_length]
    """
    return tf.strings.reduce_join(
        inputs=self.table.lookup(tf.to_int64(ids)), axis=1)
//  [python]
def get_text(self, ids):
    """Returns a string corresponding to a sequence of character ids.

        Args:
          ids: a tensor with shape [batch_size, max_sequence_length]
    """
    return tf.strings.reduce_join(
        inputs=self.table.lookup(tf.cast(ids, dtype=tf.int64)), axis=1)
//  [python]
def test_label_map():
  label_map_obj = LabelMap()
  init_op = tf.group(
    tf.global_variables_initializer(), tf.local_variables_initializer(),
    tf.tables_initializer(),
  )
  test_string_tensor = tf.constant(['test', 'value', 'discombobulated', 'Chronographs', 'Chronographs'], dtype=tf.string)
  label_tensor, label_length_tensor = label_map_obj.text_to_labels(test_string_tensor, return_lengths=True)
  max_num_step = tf.reduce_max(label_length_tensor)
  label_length_mask_tensor = tf.to_float(tf.sequence_mask(label_length_tensor, max_num_step))
  text_tensor = label_map_obj.labels_to_text(label_tensor)

  sess = tf.Session()
  sess.run(init_op)
  print('test_string', sess.run(test_string_tensor))
  print('label_tensor', sess.run(label_tensor))
  print('label_length', sess.run(label_length_tensor))
  print('label_length_mask', sess.run(label_length_mask_tensor))
  print('text_tensor', sess.run(text_tensor))
  sess.close()
//  [python]
def test_label_map():
  label_map_obj = LabelMap()
  init_op = tf.group(
    tf.global_variables_initializer(), tf.local_variables_initializer(),
    tf.tables_initializer(),
  )
  test_string_tensor = tf.constant(['test', 'value', 'discombobulated', 'Chronographs', 'Chronographs'], dtype=tf.string)
  label_tensor, label_length_tensor = label_map_obj.text_to_labels(test_string_tensor, return_lengths=True)
  max_num_step = tf.reduce_max(label_length_tensor)
  label_length_mask_tensor = tf.cast(tf.sequence_mask(label_length_tensor, max_num_step), tf.float32)
  text_tensor = label_map_obj.labels_to_text(label_tensor)

  sess = tf.Session()
  sess.run(init_op)
  print('test_string', sess.run(test_string_tensor))
  print('label_tensor', sess.run(label_tensor))
  print('label_length', sess.run(label_length_tensor))
  print('label_length_mask', sess.run(label_length_mask_tensor))
  print('text_tensor', sess.run(text_tensor))
  sess.close()
//  [python]
def visualize_images_with_bounding_boxes(images, box_outputs, step,
                                         summary_writer):
  """Records subset of evaluation images with bounding boxes."""
  if not isinstance(images, list):
    logging.warning(
        'visualize_images_with_bounding_boxes expects list of '
        'images but received type: %s and value: %s', type(images), images)
    return

  image_shape = tf.shape(images[0])
  image_height = tf.cast(image_shape[0], tf.float32)
  image_width = tf.cast(image_shape[1], tf.float32)
  normalized_boxes = box_ops.normalize_boxes(box_outputs,
                                             [image_height, image_width])

  bounding_box_color = tf.constant([[1.0, 1.0, 0.0, 1.0]])
  image_summary = tf.image.draw_bounding_boxes(
      tf.to_float(images), normalized_boxes, bounding_box_color)
  with summary_writer.as_default():
    tf.summary.image('bounding_box_summary', image_summary, step=step)
    summary_writer.flush()
//  [python]
def visualize_images_with_bounding_boxes(images, box_outputs, step,
                                         summary_writer):
  """Records subset of evaluation images with bounding boxes."""
  if not isinstance(images, list):
    logging.warning(
        'visualize_images_with_bounding_boxes expects list of '
        'images but received type: %s and value: %s', type(images), images)
    return

  image_shape = tf.shape(images[0])
  image_height = tf.cast(image_shape[0], tf.float32)
  image_width = tf.cast(image_shape[1], tf.float32)
  normalized_boxes = box_ops.normalize_boxes(box_outputs,
                                             [image_height, image_width])

  bounding_box_color = tf.constant([[1.0, 1.0, 0.0, 1.0]])
  image_summary = tf.image.draw_bounding_boxes(
      tf.cast(images, tf.float32), normalized_boxes, bounding_box_color)
  with summary_writer.as_default():
    tf.summary.image('bounding_box_summary', image_summary, step=step)
    summary_writer.flush()
//  [python]
def lookup_indexed_value(indices, row_vecs):
  """Lookup values in each row of 'row_vecs' indexed by 'indices'.

  For each sample in the batch, look up the element for the corresponding
  index.

  Args:
    indices : A tensor of shape (batch, )
    row_vecs : A tensor of shape [batch, depth]

  Returns:
    A tensor of shape (batch, ) formed by row_vecs[i, indices[i]].
  """
  gather_indices = tf.stack((tf.range(
      tf.shape(input=row_vecs)[0], dtype=tf.int32), tf.to_int32(indices)),
      axis=1)
  return tf.gather_nd(row_vecs, gather_indices)
//  [python]
def lookup_indexed_value(indices, row_vecs):
  """Lookup values in each row of 'row_vecs' indexed by 'indices'.

  For each sample in the batch, look up the element for the corresponding
  index.

  Args:
    indices : A tensor of shape (batch, )
    row_vecs : A tensor of shape [batch, depth]

  Returns:
    A tensor of shape (batch, ) formed by row_vecs[i, indices[i]].
  """
  gather_indices = tf.stack((tf.range(
      tf.shape(input=row_vecs)[0], dtype=tf.int32), tf.cast(indices, tf.int32)),
      axis=1)
  return tf.gather_nd(row_vecs, gather_indices)
//  [python]
def find_length_by_null(predicted_chars, null_code):
  """Determine sequence length by finding null_code among predicted char IDs.

  Given the char class ID for each position, compute the sequence length.
  Note that this function computes this based on the number of null_code,
  instead of the position of the first null_code.

  Args:
    predicted_chars: A tensor of [batch x seq_length] where each element stores
      the char class ID with max probability;
    null_code: an int32, character id for the NULL.

  Returns:
    A [batch, ] tensor which stores the sequence length for each sample.
  """
  return tf.reduce_sum(
      input_tensor=tf.to_int32(tf.not_equal(null_code, predicted_chars)), axis=1)
//  [python]
def find_length_by_null(predicted_chars, null_code):
  """Determine sequence length by finding null_code among predicted char IDs.

  Given the char class ID for each position, compute the sequence length.
  Note that this function computes this based on the number of null_code,
  instead of the position of the first null_code.

  Args:
    predicted_chars: A tensor of [batch x seq_length] where each element stores
      the char class ID with max probability;
    null_code: an int32, character id for the NULL.

  Returns:
    A [batch, ] tensor which stores the sequence length for each sample.
  """
  return tf.reduce_sum(
      input_tensor=tf.cast(tf.not_equal(null_code, predicted_chars), tf.int32), axis=1)
//  [python]
def _get_crop_border(border, size):
  """Get the border of cropping."""

  border = t.cast(border, t.float32)
  size = t.cast(size, t.float32)

  i = t.math.ceil(t.math.log(2.0 * border / size) / t.math.log(2.0))
  divisor = t.pow(2.0, i)
  divisor = t.clip_by_value(divisor, 1, border)
  divisor = t.cast(divisor, t.int32)

  return t.cast(border, t.int32) // divisor
//  [python]
def _get_crop_border(border, size):
  """Get the border of cropping."""

  border = tf.cast(border, tf.float32)
  size = tf.cast(size, tf.float32)

  i = tf.math.ceil(tf.math.log(2.0 * border / size) / tf.math.log(2.0))
  divisor = tf.pow(2.0, i)
  divisor = tf.clip_by_value(divisor, 1, border)
  divisor = tf.cast(divisor, tf.int32)

  return tf.cast(border, tf.int32) // divisor
//  [python]
def null_based_length_prediction(chars_log_prob, null_code):
  """Computes length and confidence of prediction based on positions of NULLs.

  Args:
    chars_log_prob: A tensor of shape [batch x seq_length x num_char_classes]
      with log probabilities of a character;
    null_code: an int32, character id for the NULL.

  Returns:
    A tuple (text_log_prob, predicted_length), where
    text_log_prob - is a tensor of the same shape as length_log_prob.
    Element #0 of the output corresponds to probability of the empty string,
    element #seq_length - is the probability of length=seq_length.
    predicted_length is a tensor with shape [batch].
  """
  predicted_chars = t.cast(
      t.argmax(input=chars_log_prob, axis=2), dtype=t.int32)
  # We do right pad to support sequences with seq_length elements.
  text_log_prob = max_char_logprob_cumsum(
      axis_pad(chars_log_prob, axis=1, after=1))
  predicted_length = find_length_by_null(predicted_chars, null_code)
  return text_log_prob, predicted_length
//  [python]
def null_based_length_prediction(chars_log_prob, null_code):
  """Computes length and confidence of prediction based on positions of NULLs.

  Args:
    chars_log_prob: A tensor of shape [batch x seq_length x num_char_classes]
      with log probabilities of a character;
    null_code: an int32, character id for the NULL.

  Returns:
    A tuple (text_log_prob, predicted_length), where
    text_log_prob - is a tensor of the same shape as length_log_prob.
    Element #0 of the output corresponds to probability of the empty string,
    element #seq_length - is the probability of length=seq_length.
    predicted_length is a tensor with shape [batch].
  """
  predicted_chars = tf.cast(
      tf.argmax(input=chars_log_prob, axis=2), dtype=tf.int32)
  # We do right pad to support sequences with seq_length elements.
  text_log_prob = max_char_logprob_cumsum(
      axis_pad(chars_log_prob, axis=1, after=1))
  predicted_length = find_length_by_null(predicted_chars, null_code)
  return text_log_prob, predicted_length
//  [python]
def prepare_opt_inputs_dict(config, input_ids, attention_mask=None, head_mask=None):
    if attention_mask is None:
        attention_mask = tf.to_int32(tf.math.not_equal(input_ids, config.pad_token_id))
    return {"input_ids": input_ids, "attention_mask": attention_mask}
//  [python]
def prepare_opt_inputs_dict(config, input_ids, attention_mask=None, head_mask=None):
    if attention_mask is None:
        attention_mask = tf.cast(tf.math.not_equal(input_ids, config.pad_token_id), tf.int8)
    return {"input_ids": input_ids, "attention_mask": attention_mask}
//  [python]
def project_color(self, proj_vert, image):
    half_size = self.img_size // 2
    vertice = tf.round(proj_vert * half_size + half_size)
    vertice = tf.cast(vertice, np.int32)
    flatten_image = tf.reshape(image, [self.batch_size, self.img_size * self.img_size, -1])
    x_coords = tf.clip_by_value(vertice[..., 1], 0, self.img_size - 1)
    y_coords = tf.clip_by_value(vertice[..., 0], 0, self.img_size - 1)
    coords = x_coords * self.img_size + y_coords
    # proj_color = tf.gather_nd(flatten_image, coords)
    proj_color = utils.batch_gather(flatten_image, coords)
    return proj_color
//  [python]
def project_color(self, proj_vert, image):
    half_size = self.img_size // 2
    vertice = tf.cast(tf.round(proj_vert * half_size + half_size), np.int32)
    flatten_image = tf.reshape(image, [self.batch_size, self.img_size * self.img_size, -1])
    x_coords = tf.clip_by_value(vertice[..., 1], 0, self.img_size - 1)
    y_coords = tf.clip_by_value(vertice[..., 0], 0, self.img_size - 1)
    coords = x_coords * self.img_size + y_coords
    # proj_color = tf.gather_nd(flatten_image, coords)
    proj_color = utils.batch_gather(flatten_image, coords)
    return proj_color
//  [python]
def char_predictions(self, chars_logit):
    """Returns confidence scores (softmax values) for predicted characters.

    Args:
      chars_logit: chars logits, a tensor with shape [batch_size x seq_length x
        num_char_classes]

    Returns:
      A tuple (ids, log_prob, scores), where:
        ids - predicted characters, a int32 tensor with shape
          [batch_size x seq_length];
        log_prob - a log probability of all characters, a float tensor with
          shape [batch_size, seq_length, num_char_classes];
        scores - corresponding confidence scores for characters, a float
        tensor
          with shape [batch_size x seq_length].
    """
    log_prob = utils.logits_to_log_prob(chars_logit)
    ids = tf.to_int32(tf.argmax(input=log_prob, axis=2),
                  name='predicted_chars')
    mask = tf.cast(
        slim.one_hot_encoding(ids, self._params.num_char_classes), tf.bool)
    all_scores = tf.nn.softmax(chars_logit)
    selected_scores = tf.boolean_mask(
        tensor=all_scores, mask=mask, name='char_scores')
    scores = tf.reshape(
        selected_scores,
        shape=(-1, self._params.seq_length),
        name='predicted_scores')
    return ids, log_prob, scores
//  [python]
def char_predictions(self, chars_logit):
    """Returns confidence scores (softmax values) for predicted characters.

    Args:
      chars_logit: chars logits, a tensor with shape [batch_size x seq_length x
        num_char_classes]

    Returns:
      A tuple (ids, log_prob, scores), where:
        ids - predicted characters, a int32 tensor with shape
          [batch_size x seq_length];
        log_prob - a log probability of all characters, a float tensor with
          shape [batch_size, seq_length, num_char_classes];
        scores - corresponding confidence scores for characters, a float
        tensor
          with shape [batch_size x seq_length].
    """
    log_prob = utils.logits_to_log_prob(chars_logit)
    ids = tf.cast(tf.argmax(input=log_prob, axis=2),
                  name='predicted_chars', dtype=tf.int32)
    mask = tf.cast(
        slim.one_hot_encoding(ids, self._params.num_char_classes), tf.bool)
    all_scores = tf.nn.softmax(chars_logit)
    selected_scores = tf.boolean_mask(
        tensor=all_scores, mask=mask, name='char_scores')
    scores = tf.reshape(
        selected_scores,
        shape=(-1, self._params.seq_length),
        name='predicted_scores')
    return ids, log_prob, scores
//  [python]
def compute_accuracy(v_xs, v_ys):
    global prediction
    y_pre = sess.run(prediction, feed_dict={xs: v_xs})
    correct_prediction = tf.equal(tf.argmax(y_pre,1), tf.argmax(v_ys,1))
    accuracy = tf.reduce_mean(tf.to_float(correct_prediction))
    result = sess.run(accuracy, feed_dict={xs: v_xs, ys: v_ys})
    return result
//  [python]
def compute_accuracy(v_xs, v_ys):
    global prediction
    y_pre = sess.run(prediction, feed_dict={xs: v_xs})
    correct_prediction = tf.equal(tf.argmax(y_pre,1), tf.argmax(v_ys,1))
    accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))
    result = sess.run(accuracy, feed_dict={xs: v_xs, ys: v_ys})
    return result
//  [python]
def sequence_loss_fn(self, chars_logits, chars_labels):
    """Loss function for char sequence.

    Depending on values of hyper parameters it applies label smoothing and can
    also ignore all null chars after the first one.

    Args:
      chars_logits: logits for predicted characters, shape=[batch_size,
        seq_length, num_char_classes];
      chars_labels: ground truth ids of characters, shape=[batch_size,
        seq_length];
      mparams: method hyper parameters.

    Returns:
      A Tensor with shape [batch_size] - the log-perplexity for each sequence.
    """
    mparams = self._mparams['sequence_loss_fn']
    with tf.compat.v1.variable_scope('sequence_loss_fn/SLF'):
      if mparams.label_smoothing > 0:
        smoothed_one_hot_labels = self.label_smoothing_regularization(
            chars_labels, mparams.label_smoothing)
        labels_list = tf.unstack(smoothed_one_hot_labels, axis=1)
      else:
        # NOTE: in case of sparse softmax we are not using one-hot
        # encoding.
        labels_list = tf.unstack(chars_labels, axis=1)

      batch_size, seq_length, _ = chars_logits.shape.as_list()
      if mparams.ignore_nulls:
        weights = tf.ones((batch_size, seq_length), dtype=tf.float32)
      else:
        # Suppose that reject character is the last in the charset.
        reject_char = tf.constant(
            self._params.num_char_classes - 1,
            shape=(batch_size, seq_length),
            dtype=tf.int64)
        known_char = tf.not_equal(chars_labels, reject_char)
        weights = tf.to_float(known_char)

      logits_list = tf.unstack(chars_logits, axis=1)
      weights_list = tf.unstack(weights, axis=1)
      loss = tf.contrib.legacy_seq2seq.sequence_loss(
          logits_list,
          labels_list,
          weights_list,
          softmax_loss_function=get_softmax_loss_fn(mparams.label_smoothing),
          average_across_timesteps=mparams.average_across_timesteps)
      tf.compat.v1.losses.add_loss(loss)
      return loss
//  [python]
def sequence_loss_fn(self, chars_logits, chars_labels):
    """Loss function for char sequence.

    Depending on values of hyper parameters it applies label smoothing and can
    also ignore all null chars after the first one.

    Args:
      chars_logits: logits for predicted characters, shape=[batch_size,
        seq_length, num_char_classes];
      chars_labels: ground truth ids of characters, shape=[batch_size,
        seq_length];
      mparams: method hyper parameters.

    Returns:
      A Tensor with shape [batch_size] - the log-perplexity for each sequence.
    """
    mparams = self._mparams['sequence_loss_fn']
    with tf.compat.v1.variable_scope('sequence_loss_fn/SLF'):
      if mparams.label_smoothing > 0:
        smoothed_one_hot_labels = self.label_smoothing_regularization(
            chars_labels, mparams.label_smoothing)
        labels_list = tf.unstack(smoothed_one_hot_labels, axis=1)
      else:
        # NOTE: in case of sparse softmax we are not using one-hot
        # encoding.
        labels_list = tf.unstack(chars_labels, axis=1)

      batch_size, seq_length, _ = chars_logits.shape.as_list()
      if mparams.ignore_nulls:
        weights = tf.ones((batch_size, seq_length), dtype=tf.float32)
      else:
        # Suppose that reject character is the last in the charset.
        reject_char = tf.constant(
            self._params.num_char_classes - 1,
            shape=(batch_size, seq_length),
            dtype=tf.int64)
        known_char = tf.not_equal(chars_labels, reject_char)
        weights = tf.cast(known_char, dtype=tf.float32)

      logits_list = tf.unstack(chars_logits, axis=1)
      weights_list = tf.unstack(weights, axis=1)
      loss = tf.contrib.legacy_seq2seq.sequence_loss(
          logits_list,
          labels_list,
          weights_list,
          softmax_loss_function=get_softmax_loss_fn(mparams.label_smoothing),
          average_across_timesteps=mparams.average_across_timesteps)
      tf.compat.v1.losses.add_loss(loss)
      return loss
//  [python]
def masked_mse_tf(preds, labels, null_val=np.nan):
    """
    Accuracy with masking.
    :param preds:
    :param labels:
    :param null_val:
    :return:
    """
    if np.isnan(null_val):
        mask = ~tf.is_nan(labels)
    else:
        mask = tf.not_equal(labels, null_val)
    mask = tf.to_float(mask)
    mask /= tf.reduce_mean(mask)
    mask = tf.where(tf.is_nan(mask), tf.zeros_like(mask), mask)
    loss = tf.square(tf.subtract(preds, labels))
    loss = loss * mask
    loss = tf.where(tf.is_nan(loss), tf.zeros_like(loss), loss)
    return tf.reduce_mean(loss)
//  [python]
def masked_mse_tf(preds, labels, null_val=np.nan):
    """
    Accuracy with masking.
    :param preds:
    :param labels:
    :param null_val:
    :return:
    """
    if np.isnan(null_val):
        mask = ~tf.is_nan(labels)
    else:
        mask = tf.not_equal(labels, null_val)
    mask = tf.cast(mask, tf.float32)
    mask /= tf.reduce_mean(mask)
    mask = tf.where(tf.is_nan(mask), tf.zeros_like(mask), mask)
    loss = tf.square(tf.subtract(preds, labels))
    loss = loss * mask
    loss = tf.where(tf.is_nan(loss), tf.zeros_like(loss), loss)
    return tf.reduce_mean(loss)
//  [python]
def masked_mae_tf(preds, labels, null_val=np.nan):
    """
    Accuracy with masking.
    :param preds:
    :param labels:
    :param null_val:
    :return:
    """
    if np.isnan(null_val):
        mask = ~tf.is_nan(labels)
    else:
        mask = tf.not_equal(labels, null_val)
    mask = tf.to_float(mask)
    mask /= tf.reduce_mean(mask)
    mask = tf.where(tf.is_nan(mask), tf.zeros_like(mask), mask)
    loss = tf.abs(tf.subtract(preds, labels))
    loss = loss * mask
    loss = tf.where(tf.is_nan(loss), tf.zeros_like(loss), loss)
    return tf.reduce_mean(loss)
//  [python]
def masked_mae_tf(preds, labels, null_val=np.nan):
    """
    Accuracy with masking.
    :param preds:
    :param labels:
    :param null_val:
    :return:
    """
    if np.isnan(null_val):
        mask = ~tf.is_nan(labels)
    else:
        mask = tf.not_equal(labels, null_val)
    mask = tf.cast(mask, tf.float32)
    mask /= tf.reduce_mean(mask)
    mask = tf.where(tf.is_nan(mask), tf.zeros_like(mask), mask)
    loss = tf.abs(tf.subtract(preds, labels))
    loss = loss * mask
    loss = tf.where(tf.is_nan(loss), tf.zeros_like(loss), loss)
    return tf.reduce_mean(loss)
//  [python]
def check_decoder_model_past_large_inputs(self, config, inputs_dict):
        model = TFOPTModel(config=config)
        input_ids = inputs_dict["input_ids"]

        input_ids = input_ids[:1, :]
        attention_mask = inputs_dict["attention_mask"][:1, :]
        self.batch_size = 1

        # first forward pass
        outputs = model(input_ids, attention_mask=attention_mask, use_cache=True)

        output, past_key_values = outputs.to_tuple()

        # create hypothetical next token and extent to next_input_ids
        next_tokens = ids_tensor((self.batch_size, 3), config.vocab_size)
        next_attn_mask = t.cast(ids_tensor((self.batch_size, 3), 2), t.int8)

        # append to next input_ids and
        next_input_ids = tf.concat([input_ids, next_tokens], axis=-1)
        next_attention_mask = tf.concat([attention_mask, next_attn_mask], axis=-1)

        output_from_no_past = model(next_input_ids, attention_mask=next_attention_mask)[0]
        output_from_past = model(next_tokens, attention_mask=next_attention_mask, past_key_values=past_key_values)[0]

        self.parent.assertEqual(next_tokens.shape[1], output_from_past.shape[1])

        # select random slice
        random_slice_idx = int(ids_tensor((1,), output_from_past.shape[-1]))
        output_from_no_past_slice = output_from_no_past[:, -3:, random_slice_idx]
        output_from_past_slice = output_from_past[:, :, random_slice_idx]

        # test that outputs are equal for slice
        tf.debugging.assert_near(output_from_past_slice, output_from_no_past_slice, rtol=1e-3)
//  [python]
def check_decoder_model_past_large_inputs(self, config, inputs_dict):
        model = TFOPTModel(config=config)
        input_ids = inputs_dict["input_ids"]

        input_ids = input_ids[:1, :]
        attention_mask = inputs_dict["attention_mask"][:1, :]
        self.batch_size = 1

        # first forward pass
        outputs = model(input_ids, attention_mask=attention_mask, use_cache=True)

        output, past_key_values = outputs.to_tuple()

        # create hypothetical next token and extent to next_input_ids
        next_tokens = ids_tensor((self.batch_size, 3), config.vocab_size)
        next_attn_mask = tf.cast(ids_tensor((self.batch_size, 3), 2), tf.int8)

        # append to next input_ids and
        next_input_ids = tf.concat([input_ids, next_tokens], axis=-1)
        next_attention_mask = tf.concat([attention_mask, next_attn_mask], axis=-1)

        output_from_no_past = model(next_input_ids, attention_mask=next_attention_mask)[0]
        output_from_past = model(next_tokens, attention_mask=next_attention_mask, past_key_values=past_key_values)[0]

        self.parent.assertEqual(next_tokens.shape[1], output_from_past.shape[1])

        # select random slice
        random_slice_idx = int(ids_tensor((1,), output_from_past.shape[-1]))
        output_from_no_past_slice = output_from_no_past[:, -3:, random_slice_idx]
        output_from_past_slice = output_from_past[:, :, random_slice_idx]

        # test that outputs are equal for slice
        tf.debugging.assert_near(output_from_past_slice, output_from_no_past_slice, rtol=1e-3)
//  [python]
def test_batch_generation(self):
        model_id = "facebook/opt-350m"

        tokenizer = GPT2Tokenizer.from_pretrained(model_id)
        model = TFOPTForCausalLM.from_pretrained(model_id)

        tokenizer.padding_side = "left"

        # use different length sentences to test batching
        sentences = [
            "Hello, my dog is a little",
            "Today, I",
        ]

        inputs = tokenizer(sentences, return_tensors="tf", padding=True)
        input_ids = inputs["input_ids"]

        outputs = model.generate(input_ids=input_ids, attention_mask=inputs["attention_mask"])

        inputs_non_padded = tokenizer(sentences[0], return_tensors="tf").input_ids
        output_non_padded = model.generate(input_ids=inputs_non_padded)

        num_paddings = inputs_non_padded.shape[-1] - tf.math.reduce_sum(
            tf.convert_to_tensor(inputs["attention_mask"][-1], dtype=tf.int64)
        )
        inputs_padded = tokenizer(sentences[1], return_tensors="tf").input_ids
        output_padded = model.generate(input_ids=inputs_padded, max_length=model.config.max_length - num_paddings)

        batch_out_sentence = tokenizer.batch_decode(outputs, skip_special_tokens=True)
        non_padded_sentence = tokenizer.decode(output_non_padded[0], skip_special_tokens=True)
        padded_sentence = tokenizer.decode(output_padded[0], skip_special_tokens=True)

        expected_output_sentence = [
            "Hello, my dog is a little bit of a dork.\nI'm a little bit",
            "Today, I was in the middle of a conversation with a friend about the",
        ]
        self.assertListEqual(expected_output_sentence, batch_out_sentence)
        self.assertListEqual(batch_out_sentence, [non_padded_sentence, padded_sentence])
//  [python]
def test_batch_generation(self):
        model_id = "facebook/opt-350m"

        tokenizer = GPT2Tokenizer.from_pretrained(model_id)
        model = TFOPTForCausalLM.from_pretrained(model_id)

        tokenizer.padding_side = "left"

        # use different length sentences to test batching
        sentences = [
            "Hello, my dog is a little",
            "Today, I",
        ]

        inputs = tokenizer(sentences, return_tensors="tf", padding=True)
        input_ids = inputs["input_ids"]

        outputs = model.generate(input_ids=input_ids, attention_mask=inputs["attention_mask"])

        inputs_non_padded = tokenizer(sentences[0], return_tensors="tf").input_ids
        output_non_padded = model.generate(input_ids=inputs_non_padded)

        num_paddings = inputs_non_padded.shape[-1] - tf.math.reduce_sum(
            tf.cast(inputs["attention_mask"][-1], tf.int64)
        )
        inputs_padded = tokenizer(sentences[1], return_tensors="tf").input_ids
        output_padded = model.generate(input_ids=inputs_padded, max_length=model.config.max_length - num_paddings)

        batch_out_sentence = tokenizer.batch_decode(outputs, skip_special_tokens=True)
        non_padded_sentence = tokenizer.decode(output_non_padded[0], skip_special_tokens=True)
        padded_sentence = tokenizer.decode(output_padded[0], skip_special_tokens=True)

        expected_output_sentence = [
            "Hello, my dog is a little bit of a dork.\nI'm a little bit",
            "Today, I was in the middle of a conversation with a friend about the",
        ]
        self.assertListEqual(expected_output_sentence, batch_out_sentence)
        self.assertListEqual(batch_out_sentence, [non_padded_sentence, padded_sentence])
//  [python]
def _split_inference_data(self, hero_data_list):
        # new add 20180912
        each_hero_data_list = []
        for hero_index in range(self.hero_num):
            this_hero_data = []
            data_feature = t.cast(hero_data_list[hero_index], t.float32)
            print(data_feature.shape)
            print(self.hero_data_split_shape[hero_index][0])
            feature = tf.reshape(data_feature, [-1, self.hero_data_split_shape[hero_index][0]])
            this_hero_data.append(feature)
            each_hero_data_list.append(this_hero_data)

        return each_hero_data_list
//  [python]
def _split_inference_data(self, hero_data_list):
        # new add 20180912
        each_hero_data_list = []
        for hero_index in range(self.hero_num):
            this_hero_data = []
            data_feature = tf.cast(hero_data_list[hero_index], tf.float32)
            print(data_feature.shape)
            print(self.hero_data_split_shape[hero_index][0])
            feature = tf.reshape(data_feature, [-1, self.hero_data_split_shape[hero_index][0]])
            this_hero_data.append(feature)
            each_hero_data_list.append(this_hero_data)

        return each_hero_data_list
//  [python]
def NormalizeImages(images, pixel_value_scale=0.5, pixel_value_offset=0.5):
  """Normalize pixel values in image.

  Output is computed as
  normalized_images = (images - pixel_value_offset) / pixel_value_scale.

  Args:
    images: `Tensor`, images to normalize.
    pixel_value_scale: float, scale.
    pixel_value_offset: float, offset.

  Returns:
    normalized_images: `Tensor`, normalized images.
  """
  images = tf.to_float(images)
  normalized_images = tf.math.divide(
      tf.subtract(images, pixel_value_offset), pixel_value_scale)
  return normalized_images
//  [python]
def NormalizeImages(images, pixel_value_scale=0.5, pixel_value_offset=0.5):
  """Normalize pixel values in image.

  Output is computed as
  normalized_images = (images - pixel_value_offset) / pixel_value_scale.

  Args:
    images: `Tensor`, images to normalize.
    pixel_value_scale: float, scale.
    pixel_value_offset: float, offset.

  Returns:
    normalized_images: `Tensor`, normalized images.
  """
  images = tf.cast(images, tf.float32)
  normalized_images = tf.math.divide(
      tf.subtract(images, pixel_value_offset), pixel_value_scale)
  return normalized_images
//  [python]
def positional_encoding(self):
        '''
        calculate the positional encoding given the window length
        :return: positional encoding (1, window_length, 1, d_model)
        '''
        angle_rads = self.get_angles(np.arange(self.window_len)[:, np.newaxis], np.arange(self.d_model)[np.newaxis, :])

        # apply sin to even indices in the array; 2i
        angle_rads[:, 0::2] = np.sin(angle_rads[:, 0::2])

        # apply cos to odd indices in the array; 2i+1
        angle_rads[:, 1::2] = np.cos(angle_rads[:, 1::2])

        pos_encoding = angle_rads[np.newaxis, ...]

        return t.cast(pos_encoding, dtype=t.float32)
//  [python]
def positional_encoding(self):
        '''
        calculate the positional encoding given the window length
        :return: positional encoding (1, window_length, 1, d_model)
        '''
        angle_rads = self.get_angles(np.arange(self.window_len)[:, np.newaxis], np.arange(self.d_model)[np.newaxis, :])

        # apply sin to even indices in the array; 2i
        angle_rads[:, 0::2] = np.sin(angle_rads[:, 0::2])

        # apply cos to odd indices in the array; 2i+1
        angle_rads[:, 1::2] = np.cos(angle_rads[:, 1::2])

        pos_encoding = angle_rads[np.newaxis, ...]

        return tf.cast(pos_encoding, dtype=tf.float32)
//  [python]
def build_graph(self, datas, global_step=None):
        datas = tf.reshape(datas, [-1, self.hero_num, self.hero_data_len])
        data_list = tf.transpose(datas, perm=[1, 0, 2])

        # new add 20180912
        each_hero_data_list = []
        for hero_index in range(self.hero_num):
            this_hero_data_list = self._split_data(tf.to_float(data_list[hero_index]), self.hero_data_split_shape[hero_index])
            each_hero_data_list.append(this_hero_data_list)
        # build network
        each_hero_fc_result_list = self._inference(each_hero_data_list)
        # calculate loss
        with tf.xla.experimental.jit_scope(self.use_xla):
            loss = self._calculate_loss(each_hero_data_list, each_hero_fc_result_list)
        # return
        return (
            loss,
            [
                loss,
                [self.each_hero_loss_list[0][1], self.each_hero_loss_list[0][2], self.each_hero_loss_list[0][3]],
                [self.each_hero_loss_list[1][1], self.each_hero_loss_list[1][2], self.each_hero_loss_list[1][3]],
                [self.each_hero_loss_list[2][1], self.each_hero_loss_list[2][2], self.each_hero_loss_list[2][3]],
            ],
            [],
        )
//  [python]
def build_graph(self, datas, global_step=None):
        datas = tf.reshape(datas, [-1, self.hero_num, self.hero_data_len])
        data_list = tf.transpose(datas, perm=[1, 0, 2])

        # new add 20180912
        each_hero_data_list = []
        for hero_index in range(self.hero_num):
            this_hero_data_list = self._split_data(tf.cast(data_list[hero_index], tf.float32), self.hero_data_split_shape[hero_index])
            each_hero_data_list.append(this_hero_data_list)
        # build network
        each_hero_fc_result_list = self._inference(each_hero_data_list)
        # calculate loss
        with tf.xla.experimental.jit_scope(self.use_xla):
            loss = self._calculate_loss(each_hero_data_list, each_hero_fc_result_list)
        # return
        return (
            loss,
            [
                loss,
                [self.each_hero_loss_list[0][1], self.each_hero_loss_list[0][2], self.each_hero_loss_list[0][3]],
                [self.each_hero_loss_list[1][1], self.each_hero_loss_list[1][2], self.each_hero_loss_list[1][3]],
                [self.each_hero_loss_list[2][1], self.each_hero_loss_list[2][2], self.each_hero_loss_list[2][3]],
            ],
            [],
        )
//  [python]
def save(self, prebatch):
        if self._batches_saved % self.batches_per_file == 0:
            self._new_file()

        if self.multihot:
            serialized = serialize_composite(t.cast(prebatch, self.dtype)).numpy()
        else:
            if isinstance(prebatch, tf.RaggedTensor):
                prebatch = prebatch.to_tensor()

            serialized = tf.io.serialize_tensor(t.cast(prebatch, dtype=self.dtype)).numpy()

        features = tf.train.Features(feature={
            self.feature_name: tf.train.Feature(bytes_list=tf.train.BytesList(value=[serialized]))
        })

        example = tf.train.Example(features=features)
        self.writer.write(example.SerializeToString())
        self._batches_saved += 1
//  [python]
def save(self, prebatch):
        if self._batches_saved % self.batches_per_file == 0:
            self._new_file()

        if self.multihot:
            serialized = serialize_composite(tf.cast(prebatch, self.dtype)).numpy()
        else:
            if isinstance(prebatch, tf.RaggedTensor):
                prebatch = prebatch.to_tensor()

            serialized = tf.io.serialize_tensor(tf.cast(prebatch, dtype=self.dtype)).numpy()

        features = tf.train.Features(feature={
            self.feature_name: tf.train.Feature(bytes_list=tf.train.BytesList(value=[serialized]))
        })

        example = tf.train.Example(features=features)
        self.writer.write(example.SerializeToString())
        self._batches_saved += 1