//  [python]
def test_exported_callables(self, use_sp_model):
    preprocess = tf.saved_model.load(
        self._do_export(
            ["d", "ef", "abc", "xy"],
            do_lower_case=True,
            # TODO(b/181866850): drop this.
            tokenize_with_offsets=not use_sp_model,
            # TODO(b/175369555): drop this.
            experimental_disable_assert=True,
            use_sp_model=use_sp_model))

    def fold_dim(rt):
      """Removes the word/subword distinction of BertTokenizer."""
      return rt if use_sp_model else rt.merge_dims(1, 2)

    # .tokenize()
    inputs = tf.constant(["abc d ef", "ABC D EF d"])
    token_ids = preprocess.tokenize(inputs)
    self.assertAllEqual(
        fold_dim(token_ids), tf.ragged.constant([[6, 4, 5], [6, 4, 5, 4]]))

    special_tokens_dict = {
        k: v.numpy().item()  # Expecting eager Tensor, converting to Python.
        for k, v in preprocess.tokenize.get_special_tokens_dict().items()
    }
    self.assertDictEqual(
        special_tokens_dict,
        dict(
            padding_id=0,
            start_of_sequence_id=2,
            end_of_segment_id=3,
            vocab_size=4 + 6 if use_sp_model else 4 + 4))

    # .tokenize_with_offsets()
    if use_sp_model:
      # TODO(b/181866850): Enable tokenize_with_offsets when it works and test.
      self.assertFalse(hasattr(preprocess, "tokenize_with_offsets"))
    else:
      token_ids, start_offsets, limit_offsets = (
          preprocess.tokenize_with_offsets(inputs))
      self.assertAllEqual(
          fold_dim(token_ids), tf.ragged.constant([[6, 4, 5], [6, 4, 5, 4]]))
      self.assertAllEqual(
          fold_dim(start_offsets), tf.ragged.constant([[0, 4, 6], [0, 4, 6,
                                                                   9]]))
      self.assertAllEqual(
          fold_dim(limit_offsets), tf.ragged.constant([[3, 5, 8], [3, 5, 8,
                                                                   10]]))
      self.assertIs(preprocess.tokenize.get_special_tokens_dict,
                    preprocess.tokenize_with_offsets.get_special_tokens_dict)

    # Root callable.
    bert_inputs = preprocess(inputs)
    self.assertAllEqual(bert_inputs["input_word_ids"].shape.as_list(), [2, 128])
    self.assertAllEqual(
        bert_inputs["input_word_ids"][:, :10],
        tf.constant([[2, 6, 4, 5, 3, 0, 0, 0, 0, 0],
                     [2, 6, 4, 5, 4, 3, 0, 0, 0, 0]]))
    self.assertAllEqual(bert_inputs["input_mask"].shape.as_list(), [2, 128])
    self.assertAllEqual(
        bert_inputs["input_mask"][:, :10],
        tf.constant([[1, 1, 1, 1, 1, 0, 0, 0, 0, 0],
                     [1, 1, 1, 1, 1, 1, 0, 0, 0, 0]]))
    self.assertAllEqual(bert_inputs["input_type_ids"].shape.as_list(), [2, 128])
    self.assertAllEqual(
        bert_inputs["input_type_ids"][:, :10],
        tf.constant([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
                     [0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]))

    # .bert_pack_inputs()
    inputs_2 = tf.constant(["d xy", "xy abc"])
    token_ids_2 = preprocess.tokenize(inputs_2)
    bert_inputs = preprocess.bert_pack_inputs([token_ids, token_ids_2],
                                              seq_length=256)
    self.assertAllEqual(bert_inputs["input_word_ids"].shape.as_list(), [2, 256])
    self.assertAllEqual(
        bert_inputs["input_word_ids"][:, :10],
        tf.constant([[2, 6, 4, 5, 3, 4, 7, 3, 0, 0],
                     [2, 6, 4, 5, 4, 3, 7, 6, 3, 0]]))
    self.assertAllEqual(bert_inputs["input_mask"].shape.as_list(), [2, 256])
    self.assertAllEqual(
        bert_inputs["input_mask"][:, :10],
        tf.constant([[1, 1, 1, 1, 1, 1, 1, 1, 0, 0],
                     [1, 1, 1, 1, 1, 1, 1, 1, 1, 0]]))
    self.assertAllEqual(bert_inputs["input_type_ids"].shape.as_list(), [2, 256])
    self.assertAllEqual(
        bert_inputs["input_type_ids"][:, :10],
        tf.constant([[0, 0, 0, 0, 0, 1, 1, 1, 0, 0],
                     [0, 0, 0, 0, 0, 0, 1, 1, 1, 0]]))