jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_reexport(self, use_sp_model):
    """Test that preprocess keeps working after another save/load cycle."""
    path1 = self._do_export(
        ["d", "ef", "abc", "xy"],
        do_lower_case=True,
        default_seq_length=10,
        tokenize_with_offsets=False,
        experimental_disable_assert=True,  # TODO(b/175369555): drop this.
        use_sp_model=use_sp_model)
    path2 = path1.rstrip("/") + ".2"
    model1 = tf.saved_model.load(path1)
    tf.saved_model.save(model1, path2)
    # Delete the first SavedModel to test that the sceond one loads by itself.
    # https://github.com/tensorflow/tensorflow/issues/46456 reports such a
    # failure case for BertTokenizer.
    tf.io.gfile.rmtree(path1)
    model2 = tf.saved_model.load(path2)

    inputs = tf.constant(["abc d ef", "ABC D EF d"])
    bert_inputs = model2(inputs)
    self.assertAllEqual(
        bert_inputs["input_word_ids"],
        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"],
        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"],
        tf.constant([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
                     [0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]))
//  [python]
def test_shapes(self, use_sp_model):
    preprocess = tf.saved_model.load(
        self._do_export(
            ["abc", "def"],
            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 expected_bert_input_shapes(batch_size, seq_length):
      return dict(
          input_word_ids=[batch_size, seq_length],
          input_mask=[batch_size, seq_length],
          input_type_ids=[batch_size, seq_length])

    for batch_size in [7, None]:
      if use_sp_model:
        token_out_shape = [batch_size, None]  # No word/subword distinction.
      else:
        token_out_shape = [batch_size, None, None]
      self.assertEqual(
          _result_shapes_in_tf_function(preprocess.tokenize,
                                        tf.TensorSpec([batch_size], tf.string)),
          token_out_shape, "with batch_size=%s" % batch_size)
      # TODO(b/181866850): Enable tokenize_with_offsets when it works and test.
      if use_sp_model:
        self.assertFalse(hasattr(preprocess, "tokenize_with_offsets"))
      else:
        self.assertEqual(
            _result_shapes_in_tf_function(
                preprocess.tokenize_with_offsets,
                tf.TensorSpec([batch_size], tf.string)), [token_out_shape] * 3,
            "with batch_size=%s" % batch_size)
      self.assertEqual(
          _result_shapes_in_tf_function(
              preprocess.bert_pack_inputs,
              [tf.RaggedTensorSpec([batch_size, None, None], tf.int32)] * 2,
              seq_length=256), expected_bert_input_shapes(batch_size, 256),
          "with batch_size=%s" % batch_size)
      self.assertEqual(
          _result_shapes_in_tf_function(preprocess,
                                        tf.TensorSpec([batch_size], tf.string)),
          expected_bert_input_shapes(batch_size, 128),
          "with batch_size=%s" % batch_size)
//  [python]
def test_special_tokens_in_estimator(self, use_sp_model):
    """Tests getting special tokens without an Eager init context."""
    preprocess_export_path = self._do_export(["d", "ef", "abc", "xy"],
                                             do_lower_case=True,
                                             use_sp_model=use_sp_model,
                                             tokenize_with_offsets=False)

    def _get_special_tokens_dict(obj):
      """Returns special tokens of restored tokenizer as Python values."""
      if tf.executing_eagerly():
        special_tokens_numpy = {
            k: v.numpy() for k, v in obj.get_special_tokens_dict()
        }
      else:
        with tf.Graph().as_default():
          # This code expects `get_special_tokens_dict()` to be a tf.function
          # with no dependencies (bound args) from the context it was loaded in,
          # and boldly assumes that it can just be called in a dfferent context.
          special_tokens_tensors = obj.get_special_tokens_dict()
          with tf.compat.v1.Session() as sess:
            special_tokens_numpy = sess.run(special_tokens_tensors)
      return {
          k: v.item()  # Numpy to Python.
          for k, v in special_tokens_numpy.items()
      }

    def input_fn():
      self.assertFalse(tf.executing_eagerly())
      # Build a preprocessing Model.
      sentences = tf_keras.layers.Input(shape=[], dtype=tf.string)
      preprocess = tf.compat.v1.saved_model.load(preprocess_export_path)
      tokenize = hub.KerasLayer(preprocess.tokenize)
      special_tokens_dict = _get_special_tokens_dict(tokenize.resolved_object)
      for k, v in special_tokens_dict.items():
        self.assertIsInstance(v, int, "Unexpected type for {}".format(k))
      tokens = tokenize(sentences)
      packed_inputs = layers.BertPackInputs(
          4, special_tokens_dict=special_tokens_dict)(
              tokens)
      preprocessing = tf_keras.Model(sentences, packed_inputs)
      # Map the dataset.
      ds = tf.data.Dataset.from_tensors(
          (tf.constant(["abc", "D EF"]), tf.constant([0, 1])))
      ds = ds.map(lambda features, labels: (preprocessing(features), labels))
      return ds

    def model_fn(features, labels, mode):
      del labels  # Unused.
      return tf_estimator.EstimatorSpec(
          mode=mode, predictions=features["input_word_ids"])

    estimator = tf_estimator.Estimator(model_fn=model_fn)
    outputs = list(estimator.predict(input_fn))
    self.assertAllEqual(outputs, np.array([[2, 6, 3, 0], [2, 4, 5, 3]]))
//  [python]
def test_special_tokens_in_estimator(self, use_sp_model):
    """Tests getting special tokens without an Eager init context."""
    preprocess_export_path = self._do_export(["d", "ef", "abc", "xy"],
                                             do_lower_case=True,
                                             use_sp_model=use_sp_model,
                                             tokenize_with_offsets=False)

    def _get_special_tokens_dict(obj):
      """Returns special tokens of restored tokenizer as Python values."""
      if tf.executing_eagerly():
        special_tokens_numpy = {
            k: v.numpy() for k, v in obj.get_special_tokens_dict()
        }
      else:
        with tf.Graph().as_default():
          # This code expects `get_special_tokens_dict()` to be a tf.function
          # with no dependencies (bound args) from the context it was loaded in,
          # and boldly assumes that it can just be called in a dfferent context.
          special_tokens_tensors = obj.get_special_tokens_dict()
          with tf.compat.v1.Session() as sess:
            special_tokens_numpy = sess.run(special_tokens_tensors)
      return {
          k: v.item()  # Numpy to Python.
          for k, v in special_tokens_numpy.items()
      }

    def input_fn():
      self.assertFalse(tf.executing_eagerly())
      # Build a preprocessing Model.
      sentences = tf_keras.layers.Input(shape=[], dtype=tf.string)
      preprocess = tf.saved_model.load(preprocess_export_path)
      tokenize = hub.KerasLayer(preprocess.tokenize)
      special_tokens_dict = _get_special_tokens_dict(tokenize.resolved_object)
      for k, v in special_tokens_dict.items():
        self.assertIsInstance(v, int, "Unexpected type for {}".format(k))
      tokens = tokenize(sentences)
      packed_inputs = layers.BertPackInputs(
          4, special_tokens_dict=special_tokens_dict)(
              tokens)
      preprocessing = tf_keras.Model(sentences, packed_inputs)
      # Map the dataset.
      ds = tf.data.Dataset.from_tensors(
          (tf.constant(["abc", "D EF"]), tf.constant([0, 1])))
      ds = ds.map(lambda features, labels: (preprocessing(features), labels))
      return ds

    def model_fn(features, labels, mode):
      del labels  # Unused.
      return tf_estimator.EstimatorSpec(
          mode=mode, predictions=features["input_word_ids"])

    estimator = tf_estimator.Estimator(model_fn=model_fn)
    outputs = list(estimator.predict(input_fn))
    self.assertAllEqual(outputs, np.array([[2, 6, 3, 0], [2, 4, 5, 3]]))
//  [python]
def dump_graph_structure(tf_dir: str, base_output: str) -> None:
    """export input and output structure of graph"""
    loaded = tf.saved_model.load(tf_dir)
    sig = loaded.signatures["serving_default"]

    # note: for input tensors, tensor name == spec name
    graph = {
        "inputs": {
            k: {
                "name": v.name,
                "shape": v.shape.as_list()
            }
            for k, v
            in sig.structured_input_signature[1].items()
        },
        "outputs": {
            k: {
                "name": v.name,
                "shape": v.shape.as_list()
            }
            for k, v
            in sig.structured_outputs.items()
        }
    }

    io_dir = os.path.join(base_output, "graph_io")
    os.makedirs(io_dir, exist_ok=False)
    io_file = os.path.join(io_dir, "graph_io.json")
    with open(io_file, "w") as f:
        json.dump(graph, f)
//  [python]
def dump_graph_structure(tf_dir: str, base_output: str) -> None:
    """export input and output structure of graph"""
    loaded = tensorflow.saved_model.load(tf_dir)
    sig = loaded.signatures["serving_default"]

    # note: for input tensors, tensor name == spec name
    graph = {
        "inputs": {
            k: {
                "name": v.name,
                "shape": v.shape.as_list()
            }
            for k, v
            in sig.structured_input_signature[1].items()
        },
        "outputs": {
            k: {
                "name": v.name,
                "shape": v.shape.as_list()
            }
            for k, v
            in sig.structured_outputs.items()
        }
    }

    io_dir = os.path.join(base_output, "graph_io")
    os.makedirs(io_dir, exist_ok=False)
    io_file = os.path.join(io_dir, "graph_io.json")
    with open(io_file, "w") as f:
        json.dump(graph, f)
//  [python]
def save_trt():

  if FLAGS.quantize_mode == 'int8':
    conversion_params = trt.DEFAULT_TRT_CONVERSION_PARAMS._replace(
      precision_mode=trt.TrtPrecisionMode.INT8,
      max_workspace_size_bytes=4000000000,
      use_calibration=True,
      max_batch_size=8)
    converter = trt.TrtGraphConverterV2(
      input_saved_model_dir=FLAGS.weights,
      conversion_params=conversion_params)
    converter.convert(calibration_input_fn=representative_data_gen)
  elif FLAGS.quantize_mode == 'float16':
    conversion_params = trt.DEFAULT_TRT_CONVERSION_PARAMS._replace(
      precision_mode=trt.TrtPrecisionMode.FP16,
      max_workspace_size_bytes=4000000000,
      max_batch_size=8)
    converter = trt.TrtGraphConverterV2(
      input_saved_model_dir=FLAGS.weights, conversion_params=conversion_params)
    converter.convert()
  else :
    conversion_params = trt.DEFAULT_TRT_CONVERSION_PARAMS._replace(
      precision_mode=trt.TrtPrecisionMode.FP32,
      max_workspace_size_bytes=4000000000,
      max_batch_size=8)
    converter = trt.TrtGraphConverterV2(
      input_saved_model_dir=FLAGS.weights, conversion_params=conversion_params)
    converter.convert()

  # converter.build(input_fn=representative_data_gen)
  converter.save(output_saved_model_dir=FLAGS.output)
  print('Done Converting to TF-TRT')

  saved_model_loaded = tf.compat.v1.saved_model.load(FLAGS.output)
  graph_func = saved_model_loaded.signatures[
    signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY]
  trt_graph = graph_func.graph.as_graph_def()
  for n in trt_graph.node:
    print(n.op)
    if n.op == "TRTEngineOp":
      print("Node: %s, %s" % (n.op, n.name.replace("/", "_")))
    else:
      print("Exclude Node: %s, %s" % (n.op, n.name.replace("/", "_")))
  logging.info("model saved to: {}".format(FLAGS.output))

  trt_engine_nodes = len([1 for n in trt_graph.node if str(n.op) == 'TRTEngineOp'])
  print("numb. of trt_engine_nodes in TensorRT graph:", trt_engine_nodes)
  all_nodes = len([1 for n in trt_graph.node])
  print("numb. of all_nodes in TensorRT graph:", all_nodes)
//  [python]
def save_trt():

  if FLAGS.quantize_mode == 'int8':
    conversion_params = trt.DEFAULT_TRT_CONVERSION_PARAMS._replace(
      precision_mode=trt.TrtPrecisionMode.INT8,
      max_workspace_size_bytes=4000000000,
      use_calibration=True,
      max_batch_size=8)
    converter = trt.TrtGraphConverterV2(
      input_saved_model_dir=FLAGS.weights,
      conversion_params=conversion_params)
    converter.convert(calibration_input_fn=representative_data_gen)
  elif FLAGS.quantize_mode == 'float16':
    conversion_params = trt.DEFAULT_TRT_CONVERSION_PARAMS._replace(
      precision_mode=trt.TrtPrecisionMode.FP16,
      max_workspace_size_bytes=4000000000,
      max_batch_size=8)
    converter = trt.TrtGraphConverterV2(
      input_saved_model_dir=FLAGS.weights, conversion_params=conversion_params)
    converter.convert()
  else :
    conversion_params = trt.DEFAULT_TRT_CONVERSION_PARAMS._replace(
      precision_mode=trt.TrtPrecisionMode.FP32,
      max_workspace_size_bytes=4000000000,
      max_batch_size=8)
    converter = trt.TrtGraphConverterV2(
      input_saved_model_dir=FLAGS.weights, conversion_params=conversion_params)
    converter.convert()

  # converter.build(input_fn=representative_data_gen)
  converter.save(output_saved_model_dir=FLAGS.output)
  print('Done Converting to TF-TRT')

  saved_model_loaded = tf.saved_model.load(FLAGS.output)
  graph_func = saved_model_loaded.signatures[
    signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY]
  trt_graph = graph_func.graph.as_graph_def()
  for n in trt_graph.node:
    print(n.op)
    if n.op == "TRTEngineOp":
      print("Node: %s, %s" % (n.op, n.name.replace("/", "_")))
    else:
      print("Exclude Node: %s, %s" % (n.op, n.name.replace("/", "_")))
  logging.info("model saved to: {}".format(FLAGS.output))

  trt_engine_nodes = len([1 for n in trt_graph.node if str(n.op) == 'TRTEngineOp'])
  print("numb. of trt_engine_nodes in TensorRT graph:", trt_engine_nodes)
  all_nodes = len([1 for n in trt_graph.node])
  print("numb. of all_nodes in TensorRT graph:", all_nodes)
//  [python]
def test_generate_tf_function_export_fixed_input_length(self):
        # TF-only test: tf.saved_model export
        test_model = TFAutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gpt2")
        input_length = 2
        max_new_tokens = 2

        class DummyModel(tf.Module):
            def __init__(self, model):
                super(DummyModel, self).__init__()
                self.model = model

            @tf.function(
                input_signature=(
                    tf.TensorSpec((None, input_length), tf.int32, name="input_ids"),
                    tf.TensorSpec((None, input_length), tf.int32, name="attention_mask"),
                ),
                jit_compile=True,
            )
            def serving(self, input_ids, attention_mask):
                outputs = self.model.generate(
                    input_ids=input_ids,
                    attention_mask=attention_mask,
                    max_new_tokens=max_new_tokens,
                    return_dict_in_generate=True,
                )
                return {"sequences": outputs["sequences"]}

        dummy_input_ids = [[2, 0], [102, 103]]
        dummy_attention_masks = [[1, 0], [1, 1]]
        dummy_model = DummyModel(model=test_model)
        with tempfile.TemporaryDirectory() as tmp_dir:
            tf.saved_model.save(dummy_model, tmp_dir, signatures={"serving_default": dummy_model.serving})
            serving_func = tf.compat.v1.saved_model.load(tmp_dir).signatures["serving_default"]
            for batch_size in range(1, len(dummy_input_ids) + 1):
                inputs = {
                    "input_ids": tf.constant(dummy_input_ids[:batch_size]),
                    "attention_mask": tf.constant(dummy_attention_masks[:batch_size]),
                }
                tf_func_outputs = serving_func(**inputs)["sequences"]
                tf_model_outputs = test_model.generate(**inputs, max_new_tokens=max_new_tokens)
                tf.debugging.assert_equal(tf_func_outputs, tf_model_outputs)
//  [python]
def test_generate_tf_function_export_fixed_input_length(self):
        # TF-only test: tf.saved_model export
        test_model = TFAutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gpt2")
        input_length = 2
        max_new_tokens = 2

        class DummyModel(tf.Module):
            def __init__(self, model):
                super(DummyModel, self).__init__()
                self.model = model

            @tf.function(
                input_signature=(
                    tf.TensorSpec((None, input_length), tf.int32, name="input_ids"),
                    tf.TensorSpec((None, input_length), tf.int32, name="attention_mask"),
                ),
                jit_compile=True,
            )
            def serving(self, input_ids, attention_mask):
                outputs = self.model.generate(
                    input_ids=input_ids,
                    attention_mask=attention_mask,
                    max_new_tokens=max_new_tokens,
                    return_dict_in_generate=True,
                )
                return {"sequences": outputs["sequences"]}

        dummy_input_ids = [[2, 0], [102, 103]]
        dummy_attention_masks = [[1, 0], [1, 1]]
        dummy_model = DummyModel(model=test_model)
        with tempfile.TemporaryDirectory() as tmp_dir:
            tf.saved_model.save(dummy_model, tmp_dir, signatures={"serving_default": dummy_model.serving})
            serving_func = tf.saved_model.load(tmp_dir).signatures["serving_default"]
            for batch_size in range(1, len(dummy_input_ids) + 1):
                inputs = {
                    "input_ids": tf.constant(dummy_input_ids[:batch_size]),
                    "attention_mask": tf.constant(dummy_attention_masks[:batch_size]),
                }
                tf_func_outputs = serving_func(**inputs)["sequences"]
                tf_model_outputs = test_model.generate(**inputs, max_new_tokens=max_new_tokens)
                tf.debugging.assert_equal(tf_func_outputs, tf_model_outputs)
//  [python]
def test_generate_tf_function_export_fixed_batch_size(self):
        # TF-only test: tf.saved_model export
        test_model = TFAutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gpt2")
        batch_size = 1
        max_new_tokens = 2

        class DummyModel(tf.Module):
            def __init__(self, model):
                super(DummyModel, self).__init__()
                self.model = model

            @tf.function(
                input_signature=(
                    tf.TensorSpec((batch_size, None), tf.int32, name="input_ids"),
                    tf.TensorSpec((batch_size, None), tf.int32, name="attention_mask"),
                ),
                jit_compile=True,
            )
            def serving(self, input_ids, attention_mask):
                outputs = self.model.generate(
                    input_ids=input_ids,
                    attention_mask=attention_mask,
                    max_new_tokens=max_new_tokens,
                    return_dict_in_generate=True,
                )
                return {"sequences": outputs["sequences"]}

        dummy_input_ids = [[2], [102, 103]]
        dummy_attention_masks = [[1], [1, 1]]
        dummy_model = DummyModel(model=test_model)
        with tempfile.TemporaryDirectory() as tmp_dir:
            tf.saved_model.save(dummy_model, tmp_dir, signatures={"serving_default": dummy_model.serving})
            serving_func = tf.compat.v1.saved_model.load(tmp_dir).signatures["serving_default"]
            for input_row in range(len(dummy_input_ids)):
                inputs = {
                    "input_ids": tf.constant([dummy_input_ids[input_row]]),
                    "attention_mask": tf.constant([dummy_attention_masks[input_row]]),
                }
                tf_func_outputs = serving_func(**inputs)["sequences"]
                tf_model_outputs = test_model.generate(**inputs, max_new_tokens=max_new_tokens)
                tf.debugging.assert_equal(tf_func_outputs, tf_model_outputs)
//  [python]
def test_generate_tf_function_export_fixed_batch_size(self):
        # TF-only test: tf.saved_model export
        test_model = TFAutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gpt2")
        batch_size = 1
        max_new_tokens = 2

        class DummyModel(tf.Module):
            def __init__(self, model):
                super(DummyModel, self).__init__()
                self.model = model

            @tf.function(
                input_signature=(
                    tf.TensorSpec((batch_size, None), tf.int32, name="input_ids"),
                    tf.TensorSpec((batch_size, None), tf.int32, name="attention_mask"),
                ),
                jit_compile=True,
            )
            def serving(self, input_ids, attention_mask):
                outputs = self.model.generate(
                    input_ids=input_ids,
                    attention_mask=attention_mask,
                    max_new_tokens=max_new_tokens,
                    return_dict_in_generate=True,
                )
                return {"sequences": outputs["sequences"]}

        dummy_input_ids = [[2], [102, 103]]
        dummy_attention_masks = [[1], [1, 1]]
        dummy_model = DummyModel(model=test_model)
        with tempfile.TemporaryDirectory() as tmp_dir:
            tf.saved_model.save(dummy_model, tmp_dir, signatures={"serving_default": dummy_model.serving})
            serving_func = tf.saved_model.load(tmp_dir).signatures["serving_default"]
            for input_row in range(len(dummy_input_ids)):
                inputs = {
                    "input_ids": tf.constant([dummy_input_ids[input_row]]),
                    "attention_mask": tf.constant([dummy_attention_masks[input_row]]),
                }
                tf_func_outputs = serving_func(**inputs)["sequences"]
                tf_model_outputs = test_model.generate(**inputs, max_new_tokens=max_new_tokens)
                tf.debugging.assert_equal(tf_func_outputs, tf_model_outputs)
//  [python]
def model(self):
        """Return model in AutoTrackable object."""
        if self._auto_trackable:
            return self._auto_trackable

        root = os.path.abspath(os.path.expanduser(cfg.default_workspace))
        root += str(time.time())
        if os.path.exists(root):
            shutil.rmtree(root)
        os.makedirs(root, exist_ok=True)
        if not self._sess:
            self._load_sess(self._model, **self.kwargs)
        _, builder = self.build_saved_model(root)
        builder.save()
        model = tf.compat.v1.saved_model.load(root)
        shutil.rmtree(root)
        self._auto_trackable = model
        return model
//  [python]
def model(self):
        """Return model in AutoTrackable object."""
        if self._auto_trackable:
            return self._auto_trackable

        root = os.path.abspath(os.path.expanduser(cfg.default_workspace))
        root += str(time.time())
        if os.path.exists(root):
            shutil.rmtree(root)
        os.makedirs(root, exist_ok=True)
        if not self._sess:
            self._load_sess(self._model, **self.kwargs)
        _, builder = self.build_saved_model(root)
        builder.save()
        model = tf.saved_model.load(root)
        shutil.rmtree(root)
        self._auto_trackable = model
        return model
//  [python]
def model(self):
        """Return model in AutoTrackable Format."""
        if self._sq_weight_scale_dict:
            self.adjust_weight(self.graph_def)
        if not self._auto_trackable:
            self._auto_trackable = tf.compat.v1.saved_model.load(self._model)
        return self._auto_trackable
//  [python]
def model(self):
        """Return model in AutoTrackable Format."""
        if self._sq_weight_scale_dict:
            self.adjust_weight(self.graph_def)
        if not self._auto_trackable:
            self._auto_trackable = tf.saved_model.load(self._model)
        return self._auto_trackable
//  [python]
def test_special_tokens_in_estimator(self, use_sp_model):
    """Tests getting special tokens without an Eager init context."""
    preprocess_export_path = self._do_export(["d", "ef", "abc", "xy"],
                                             do_lower_case=True,
                                             use_sp_model=use_sp_model,
                                             tokenize_with_offsets=False)

    def _get_special_tokens_dict(obj):
      """Returns special tokens of restored tokenizer as Python values."""
      if tf.executing_eagerly():
        special_tokens_numpy = {
            k: v.numpy() for k, v in obj.get_special_tokens_dict()
        }
      else:
        with tf.Graph().as_default():
          # This code expects `get_special_tokens_dict()` to be a tf.function
          # with no dependencies (bound args) from the context it was loaded in,
          # and boldly assumes that it can just be called in a dfferent context.
          special_tokens_tensors = obj.get_special_tokens_dict()
          with tf.compat.v1.Session() as sess:
            special_tokens_numpy = sess.run(special_tokens_tensors)
      return {
          k: v.item()  # Numpy to Python.
          for k, v in special_tokens_numpy.items()
      }

    def input_fn():
      self.assertFalse(tf.executing_eagerly())
      # Build a preprocessing Model.
      sentences = tf.keras.layers.Input(shape=[], dtype=tf.string)
      preprocess = tf.compat.v1.saved_model.load(preprocess_export_path)
      tokenize = hub.KerasLayer(preprocess.tokenize)
      special_tokens_dict = _get_special_tokens_dict(tokenize.resolved_object)
      for k, v in special_tokens_dict.items():
        self.assertIsInstance(v, int, "Unexpected type for {}".format(k))
      tokens = tokenize(sentences)
      packed_inputs = layers.BertPackInputs(
          4, special_tokens_dict=special_tokens_dict)(
              tokens)
      preprocessing = tf.keras.Model(sentences, packed_inputs)
      # Map the dataset.
      ds = tf.data.Dataset.from_tensors(
          (tf.constant(["abc", "D EF"]), tf.constant([0, 1])))
      ds = ds.map(lambda features, labels: (preprocessing(features), labels))
      return ds

    def model_fn(features, labels, mode):
      del labels  # Unused.
      return tf_estimator.EstimatorSpec(
          mode=mode, predictions=features["input_word_ids"])

    estimator = tf_estimator.Estimator(model_fn=model_fn)
    outputs = list(estimator.predict(input_fn))
    self.assertAllEqual(outputs, np.array([[2, 6, 3, 0], [2, 4, 5, 3]]))
//  [python]
def test_special_tokens_in_estimator(self, use_sp_model):
    """Tests getting special tokens without an Eager init context."""
    preprocess_export_path = self._do_export(["d", "ef", "abc", "xy"],
                                             do_lower_case=True,
                                             use_sp_model=use_sp_model,
                                             tokenize_with_offsets=False)

    def _get_special_tokens_dict(obj):
      """Returns special tokens of restored tokenizer as Python values."""
      if tf.executing_eagerly():
        special_tokens_numpy = {
            k: v.numpy() for k, v in obj.get_special_tokens_dict()
        }
      else:
        with tf.Graph().as_default():
          # This code expects `get_special_tokens_dict()` to be a tf.function
          # with no dependencies (bound args) from the context it was loaded in,
          # and boldly assumes that it can just be called in a dfferent context.
          special_tokens_tensors = obj.get_special_tokens_dict()
          with tf.compat.v1.Session() as sess:
            special_tokens_numpy = sess.run(special_tokens_tensors)
      return {
          k: v.item()  # Numpy to Python.
          for k, v in special_tokens_numpy.items()
      }

    def input_fn():
      self.assertFalse(tf.executing_eagerly())
      # Build a preprocessing Model.
      sentences = tf.keras.layers.Input(shape=[], dtype=tf.string)
      preprocess = tf.saved_model.load(preprocess_export_path)
      tokenize = hub.KerasLayer(preprocess.tokenize)
      special_tokens_dict = _get_special_tokens_dict(tokenize.resolved_object)
      for k, v in special_tokens_dict.items():
        self.assertIsInstance(v, int, "Unexpected type for {}".format(k))
      tokens = tokenize(sentences)
      packed_inputs = layers.BertPackInputs(
          4, special_tokens_dict=special_tokens_dict)(
              tokens)
      preprocessing = tf.keras.Model(sentences, packed_inputs)
      # Map the dataset.
      ds = tf.data.Dataset.from_tensors(
          (tf.constant(["abc", "D EF"]), tf.constant([0, 1])))
      ds = ds.map(lambda features, labels: (preprocessing(features), labels))
      return ds

    def model_fn(features, labels, mode):
      del labels  # Unused.
      return tf_estimator.EstimatorSpec(
          mode=mode, predictions=features["input_word_ids"])

    estimator = tf_estimator.Estimator(model_fn=model_fn)
    outputs = list(estimator.predict(input_fn))
    self.assertAllEqual(outputs, np.array([[2, 6, 3, 0], [2, 4, 5, 3]]))
//  [python]
def test_generate_tf_function_export(self):
        test_model = TFAutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gpt2")
        max_length = 2

        class DummyModel(tf.Module):
            def __init__(self, model):
                super(DummyModel, self).__init__()
                self.model = model

            @tf.function(
                input_signature=(
                    tf.TensorSpec((None, max_length), tf.int32, name="input_ids"),
                    tf.TensorSpec((None, max_length), tf.int32, name="attention_mask"),
                ),
                jit_compile=True,
            )
            def serving(self, input_ids, attention_mask):
                outputs = self.model.generate(
                    input_ids=input_ids,
                    attention_mask=attention_mask,
                    max_new_tokens=max_length,
                    return_dict_in_generate=True,
                )
                return {"sequences": outputs["sequences"]}

        dummy_input_ids = [[2, 0], [102, 103]]
        dummy_attention_masks = [[1, 0], [1, 1]]
        dummy_model = DummyModel(model=test_model)
        with tempfile.TemporaryDirectory() as tmp_dir:
            tf.saved_model.save(dummy_model, tmp_dir, signatures={"serving_default": dummy_model.serving})
            serving_func = tf.compat.v1.saved_model.load(tmp_dir).signatures["serving_default"]
            for batch_size in range(1, len(dummy_input_ids) + 1):
                inputs = {
                    "input_ids": tf.constant(dummy_input_ids[:batch_size]),
                    "attention_mask": tf.constant(dummy_attention_masks[:batch_size]),
                }
                tf_func_outputs = serving_func(**inputs)["sequences"]
                tf_model_outputs = test_model.generate(**inputs, max_new_tokens=max_length)
                tf.debugging.assert_equal(tf_func_outputs, tf_model_outputs)
//  [python]
def test_generate_tf_function_export(self):
        test_model = TFAutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gpt2")
        max_length = 2

        class DummyModel(tf.Module):
            def __init__(self, model):
                super(DummyModel, self).__init__()
                self.model = model

            @tf.function(
                input_signature=(
                    tf.TensorSpec((None, max_length), tf.int32, name="input_ids"),
                    tf.TensorSpec((None, max_length), tf.int32, name="attention_mask"),
                ),
                jit_compile=True,
            )
            def serving(self, input_ids, attention_mask):
                outputs = self.model.generate(
                    input_ids=input_ids,
                    attention_mask=attention_mask,
                    max_new_tokens=max_length,
                    return_dict_in_generate=True,
                )
                return {"sequences": outputs["sequences"]}

        dummy_input_ids = [[2, 0], [102, 103]]
        dummy_attention_masks = [[1, 0], [1, 1]]
        dummy_model = DummyModel(model=test_model)
        with tempfile.TemporaryDirectory() as tmp_dir:
            tf.saved_model.save(dummy_model, tmp_dir, signatures={"serving_default": dummy_model.serving})
            serving_func = tf.saved_model.load(tmp_dir).signatures["serving_default"]
            for batch_size in range(1, len(dummy_input_ids) + 1):
                inputs = {
                    "input_ids": tf.constant(dummy_input_ids[:batch_size]),
                    "attention_mask": tf.constant(dummy_attention_masks[:batch_size]),
                }
                tf_func_outputs = serving_func(**inputs)["sequences"]
                tf_model_outputs = test_model.generate(**inputs, max_new_tokens=max_length)
                tf.debugging.assert_equal(tf_func_outputs, tf_model_outputs)
//  [python]
def __init__(self, model_dir, precision, output_tensor_name="output_1"):
        temp_tftrt_dir = f"/tmp/tf-trt_model_{precision}"
        export_model(model_dir, precision, temp_tftrt_dir)
        saved_model_loaded = tf.compat.v1.saved_model.load(temp_tftrt_dir, tags=[tag_constants.SERVING])
        print(f"TF-TRT model loaded from {temp_tftrt_dir}")
        self.graph_func = saved_model_loaded.signatures[signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY]
        self.output_tensor_name = output_tensor_name
        self.precision = tf.float16 if precision == "amp" else tf.float32
//  [python]
def __init__(self, model_dir, precision, output_tensor_name="output_1"):
        temp_tftrt_dir = f"/tmp/tf-trt_model_{precision}"
        export_model(model_dir, precision, temp_tftrt_dir)
        saved_model_loaded = tf.saved_model.load(temp_tftrt_dir, tags=[tag_constants.SERVING])
        print(f"TF-TRT model loaded from {temp_tftrt_dir}")
        self.graph_func = saved_model_loaded.signatures[signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY]
        self.output_tensor_name = output_tensor_name
        self.precision = tf.float16 if precision == "amp" else tf.float32
//  [python]
def test_reexport(self, use_sp_model):
    """Test that preprocess keeps working after another save/load cycle."""
    path1 = self._do_export(
        ["d", "ef", "abc", "xy"], do_lower_case=True, default_seq_length=10,
        tokenize_with_offsets=False,
        experimental_disable_assert=True,  # TODO(b/175369555): drop this.
        use_sp_model=use_sp_model)
    path2 = path1.rstrip("/") + ".2"
    model1 = tf.saved_model.load(path1)
    tf.saved_model.save(model1, path2)
    # Delete the first SavedModel to test that the sceond one loads by itself.
    # https://github.com/tensorflow/tensorflow/issues/46456 reports such a
    # failure case for BertTokenizer.
    tf.io.gfile.rmtree(path1)
    model2 = tf.saved_model.load(path2)

    inputs = tf.constant(["abc d ef", "ABC D EF d"])
    bert_inputs = model2(inputs)
    self.assertAllEqual(bert_inputs["input_word_ids"],
                        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"],
                        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"],
                        tf.constant([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
                                     [0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]))
//  [python]
def test_special_tokens_in_estimator(self, use_sp_model):
    """Tests getting special tokens without an Eager init context."""
    preprocess_export_path = self._do_export(
        ["d", "ef", "abc", "xy"], do_lower_case=True,
        use_sp_model=use_sp_model, tokenize_with_offsets=False)

    def _get_special_tokens_dict(obj):
      """Returns special tokens of restored tokenizer as Python values."""
      if tf.executing_eagerly():
        special_tokens_numpy = {k: v.numpy()
                                for k, v in obj.get_special_tokens_dict()}
      else:
        with tf.Graph().as_default():
          # This code expects `get_special_tokens_dict()` to be a tf.function
          # with no dependencies (bound args) from the context it was loaded in,
          # and boldly assumes that it can just be called in a dfferent context.
          special_tokens_tensors = obj.get_special_tokens_dict()
          with tf.compat.v1.Session() as sess:
            special_tokens_numpy = sess.run(special_tokens_tensors)
      return {k: v.item()  # Numpy to Python.
              for k, v in special_tokens_numpy.items()}

    def input_fn():
      self.assertFalse(tf.executing_eagerly())
      # Build a preprocessing Model.
      sentences = tf.keras.layers.Input(shape=[], dtype=tf.string)
      preprocess = tf.compat.v1.saved_model.load(preprocess_export_path)
      tokenize = hub.KerasLayer(preprocess.tokenize)
      special_tokens_dict = _get_special_tokens_dict(tokenize.resolved_object)
      for k, v in special_tokens_dict.items():
        self.assertIsInstance(v, int, "Unexpected type for {}".format(k))
      tokens = tokenize(sentences)
      packed_inputs = layers.BertPackInputs(
          4, special_tokens_dict=special_tokens_dict)(tokens)
      preprocessing = tf.keras.Model(sentences, packed_inputs)
      # Map the dataset.
      ds = tf.data.Dataset.from_tensors(
          (tf.constant(["abc", "D EF"]), tf.constant([0, 1])))
      ds = ds.map(lambda features, labels: (preprocessing(features), labels))
      return ds

    def model_fn(features, labels, mode):
      del labels  # Unused.
      return tf.estimator.EstimatorSpec(mode=mode,
                                        predictions=features["input_word_ids"])

    estimator = tf.estimator.Estimator(model_fn=model_fn)
    outputs = list(estimator.predict(input_fn))
    self.assertAllEqual(outputs, np.array([[2, 6, 3, 0],
                                           [2, 4, 5, 3]]))
//  [python]
def test_special_tokens_in_estimator(self, use_sp_model):
    """Tests getting special tokens without an Eager init context."""
    preprocess_export_path = self._do_export(
        ["d", "ef", "abc", "xy"], do_lower_case=True,
        use_sp_model=use_sp_model, tokenize_with_offsets=False)

    def _get_special_tokens_dict(obj):
      """Returns special tokens of restored tokenizer as Python values."""
      if tf.executing_eagerly():
        special_tokens_numpy = {k: v.numpy()
                                for k, v in obj.get_special_tokens_dict()}
      else:
        with tf.Graph().as_default():
          # This code expects `get_special_tokens_dict()` to be a tf.function
          # with no dependencies (bound args) from the context it was loaded in,
          # and boldly assumes that it can just be called in a dfferent context.
          special_tokens_tensors = obj.get_special_tokens_dict()
          with tf.compat.v1.Session() as sess:
            special_tokens_numpy = sess.run(special_tokens_tensors)
      return {k: v.item()  # Numpy to Python.
              for k, v in special_tokens_numpy.items()}

    def input_fn():
      self.assertFalse(tf.executing_eagerly())
      # Build a preprocessing Model.
      sentences = tf.keras.layers.Input(shape=[], dtype=tf.string)
      preprocess = tf.saved_model.load(preprocess_export_path)
      tokenize = hub.KerasLayer(preprocess.tokenize)
      special_tokens_dict = _get_special_tokens_dict(tokenize.resolved_object)
      for k, v in special_tokens_dict.items():
        self.assertIsInstance(v, int, "Unexpected type for {}".format(k))
      tokens = tokenize(sentences)
      packed_inputs = layers.BertPackInputs(
          4, special_tokens_dict=special_tokens_dict)(tokens)
      preprocessing = tf.keras.Model(sentences, packed_inputs)
      # Map the dataset.
      ds = tf.data.Dataset.from_tensors(
          (tf.constant(["abc", "D EF"]), tf.constant([0, 1])))
      ds = ds.map(lambda features, labels: (preprocessing(features), labels))
      return ds

    def model_fn(features, labels, mode):
      del labels  # Unused.
      return tf.estimator.EstimatorSpec(mode=mode,
                                        predictions=features["input_word_ids"])

    estimator = tf.estimator.Estimator(model_fn=model_fn)
    outputs = list(estimator.predict(input_fn))
    self.assertAllEqual(outputs, np.array([[2, 6, 3, 0],
                                           [2, 4, 5, 3]]))
//  [python]
def test_export_module(self):
    tmp_dir = self.get_temp_dir()
    model = tf_keras.layers.Dense(2)
    inputs = tf.ones([2, 4], tf.float32)
    expected_output = model(inputs, training=False)
    module = TestModule(params=None, model=model)
    ckpt_path = tf.train.Checkpoint(model=model).save(
        os.path.join(tmp_dir, 'ckpt'))
    export_dir = export_base.export(
        module, ['foo'],
        export_savedmodel_dir=tmp_dir,
        checkpoint_path=ckpt_path,
        timestamped=True)
    self.assertTrue(os.path.exists(os.path.join(export_dir, 'saved_model.pb')))
    self.assertTrue(
        os.path.exists(
            os.path.join(export_dir, 'variables', 'variables.index')))
    self.assertTrue(
        os.path.exists(
            os.path.join(export_dir, 'variables',
                         'variables.data-00000-of-00001')))

    imported = tf.compat.v1.saved_model.load(export_dir)
    output = imported.signatures['foo'](inputs)
    self.assertAllClose(output['outputs'].numpy(), expected_output.numpy())
//  [python]
def test_export_module(self):
    tmp_dir = self.get_temp_dir()
    model = tf_keras.layers.Dense(2)
    inputs = tf.ones([2, 4], tf.float32)
    expected_output = model(inputs, training=False)
    module = TestModule(params=None, model=model)
    ckpt_path = tf.train.Checkpoint(model=model).save(
        os.path.join(tmp_dir, 'ckpt'))
    export_dir = export_base.export(
        module, ['foo'],
        export_savedmodel_dir=tmp_dir,
        checkpoint_path=ckpt_path,
        timestamped=True)
    self.assertTrue(os.path.exists(os.path.join(export_dir, 'saved_model.pb')))
    self.assertTrue(
        os.path.exists(
            os.path.join(export_dir, 'variables', 'variables.index')))
    self.assertTrue(
        os.path.exists(
            os.path.join(export_dir, 'variables',
                         'variables.data-00000-of-00001')))

    imported = tf.saved_model.load(export_dir)
    output = imported.signatures['foo'](inputs)
    self.assertAllClose(output['outputs'].numpy(), expected_output.numpy())
//  [python]
def test_custom_inference_step(self):
    tmp_dir = self.get_temp_dir()
    model = tf_keras.layers.Dense(2)
    inputs = tf.ones([2, 4], tf.float32)

    def _inference_step(inputs, model):
      return tf.nn.softmax(model(inputs, training=False))

    module = TestModule(
        params=None, model=model, inference_step=_inference_step)
    expected_output = _inference_step(inputs, model)
    ckpt_path = tf.train.Checkpoint(model=model).save(
        os.path.join(tmp_dir, 'ckpt'))
    export_dir = export_base.export(
        module, ['foo'],
        export_savedmodel_dir=tmp_dir,
        checkpoint_path=ckpt_path,
        timestamped=False)
    imported = tf.compat.v1.saved_model.load(export_dir)
    output = imported.signatures['foo'](inputs)
    self.assertAllClose(output['outputs'].numpy(), expected_output.numpy())
//  [python]
def test_custom_inference_step(self):
    tmp_dir = self.get_temp_dir()
    model = tf_keras.layers.Dense(2)
    inputs = tf.ones([2, 4], tf.float32)

    def _inference_step(inputs, model):
      return tf.nn.softmax(model(inputs, training=False))

    module = TestModule(
        params=None, model=model, inference_step=_inference_step)
    expected_output = _inference_step(inputs, model)
    ckpt_path = tf.train.Checkpoint(model=model).save(
        os.path.join(tmp_dir, 'ckpt'))
    export_dir = export_base.export(
        module, ['foo'],
        export_savedmodel_dir=tmp_dir,
        checkpoint_path=ckpt_path,
        timestamped=False)
    imported = tf.saved_model.load(export_dir)
    output = imported.signatures['foo'](inputs)
    self.assertAllClose(output['outputs'].numpy(), expected_output.numpy())
//  [python]
def _load_pyfunc(path):
    """
    Load PyFunc implementation. Called by ``pyfunc.load_pyfunc``. This function loads an MLflow
    model with the TensorFlow flavor into a new TensorFlow graph and exposes it behind the
    ``pyfunc.predict`` interface.

    :param path: Local filesystem path to the MLflow Model with the ``tensorflow`` flavor.
    """
    tf_saved_model_dir, tf_meta_graph_tags, tf_signature_def_key =\
        _get_and_parse_flavor_configuration(model_path=path)
    if LooseVersion(tensorflow.__version__) < LooseVersion('2.0.0'):
        tf_graph = tensorflow.Graph()
        tf_sess = tensorflow.Session(graph=tf_graph)
        with tf_graph.as_default():
            signature_def = _load_tensorflow_saved_model(
                tf_saved_model_dir=tf_saved_model_dir, tf_sess=tf_sess,
                tf_meta_graph_tags=tf_meta_graph_tags, tf_signature_def_key=tf_signature_def_key)

        return _TFWrapper(tf_sess=tf_sess, tf_graph=tf_graph, signature_def=signature_def)
    else:
        with tensorflow.compat.v1.Session(graph=tensorflow.Graph()) as tf_sess:
            tf.saved_model.loader.load(tf_sess, tf_meta_graph_tags, tf_saved_model_dir)
            graph = tf_sess.graph
            signature_def = graph.get_collection(tf.saved_model.signature_def_key)[0]
            return _TF2Wrapper(infer = signature_def)
//  [python]
def _load_pyfunc(path):
    """
    Load PyFunc implementation. Called by ``pyfunc.load_pyfunc``. This function loads an MLflow
    model with the TensorFlow flavor into a new TensorFlow graph and exposes it behind the
    ``pyfunc.predict`` interface.

    :param path: Local filesystem path to the MLflow Model with the ``tensorflow`` flavor.
    """
    tf_saved_model_dir, tf_meta_graph_tags, tf_signature_def_key =\
        _get_and_parse_flavor_configuration(model_path=path)
    if LooseVersion(tensorflow.__version__) < LooseVersion('2.0.0'):
        tf_graph = tensorflow.Graph()
        tf_sess = tensorflow.Session(graph=tf_graph)
        with tf_graph.as_default():
            signature_def = _load_tensorflow_saved_model(
                tf_saved_model_dir=tf_saved_model_dir, tf_sess=tf_sess,
                tf_meta_graph_tags=tf_meta_graph_tags, tf_signature_def_key=tf_signature_def_key)

        return _TFWrapper(tf_sess=tf_sess, tf_graph=tf_graph, signature_def=signature_def)
    else:
        loaded_model = tensorflow.saved_model.load(  # pylint: disable=no-value-for-parameter
                                                   export_dir=tf_saved_model_dir,
                                                   tags=tf_meta_graph_tags)
        return _TF2Wrapper(infer=loaded_model.signatures[tf_signature_def_key])
//  [python]
def initialize(self):
    logging.info("Loading model.")
    imported = tf.compat.v1.saved_model.load(self.checkpoint)
    self._bleurt_model_ops = imported.signatures["serving_default"]
//  [python]
def initialize(self):
    logging.info("Loading model.")
    imported = tf.saved_model.load(self.checkpoint)
    self._bleurt_model_ops = imported.signatures["serving_default"]