code: def save_tf_categorical_model(tmp_path):
      path = os.path.join(
          args.mlflow_repo_path,
          "tests/datasets/uci-autos-imports-85.data",
      )
      # Order is important for the csv-readers, so we use an OrderedDict here
      defaults = collections.OrderedDict(
          [
              ("body-style", [""]),
              ("curb-weight", [0.0]),
              ("highway-mpg", [0.0]),
              ("price", [0.0]),
          ]
      )
      types = collections.OrderedDict(
          (key, type(value[0])) for key, value in defaults.items()
      )
      df = pd.read_csv(
          path, names=list(types.keys()), dtype=types, na_values="?"
      )
      df = df.dropna()
  
      # Extract the label from the features dataframe
      y_train = df.pop("price")
  
      # Create the required input training function
      training_features = df.to_dict(orient="series")
  
      # Create the feature columns required for the DNNRegressor
      body_style_vocab = [
          "hardtop",
          "wagon",
          "sedan",
          "hatchback",
          "convertible",
      ]
      body_style = (
          tf.feature_column.categorical_column_with_vocabulary_list(
              key="body-style", vocabulary_list=body_style_vocab
          )
      )
      feature_columns = [
          tf.feature_column.numeric_column(key="curb-weight"),
          tf.feature_column.numeric_column(key="highway-mpg"),
          # Since this is a DNN model, convert categorical columns from sparse to dense.
          # Then, wrap them in an `indicator_column` to create a one-hot vector from the input
          tf.feature_column.indicator_column(body_style),
      ]
  
      # Build a DNNRegressor, with 20x20-unit hidden layers, with the feature columns
      # defined above as input
      estimator = tf_estimator.DNNRegressor(
          hidden_units=[20, 20], feature_columns=feature_columns
      )
  
      # Train the estimator and obtain expected predictions on the training dataset
      estimator.train(
          input_fn=lambda: iris_data_utils.train_input_fn(
              training_features, y_train, 1
          ),
          steps=10,
      )
      estimator_preds = np.array(
          [
              s["predictions"]
              for s in estimator.predict(
                  lambda: iris_data_utils.eval_input_fn(
                      training_features, None, 1
                  )
              )
          ]
      ).ravel()
      estimator_preds_df = pd.DataFrame(
          {"predictions": estimator_preds}
      )
  
      # Define a function for estimator inference
      feature_spec = {
          "body-style": tf.Variable(
              [], dtype=tf.string, name="body-style"
          ),
          "curb-weight": tf.Variable(
              [], dtype=tf.float64, name="curb-weight"
          ),
          "highway-mpg": tf.Variable(
              [], dtype=tf.float64, name="highway-mpg"
          ),
      }
      receiver_fn = (
          tf_estimator.export.build_raw_serving_input_receiver_fn(
              feature_spec
          )
      )
  
      # Save the estimator and its inference function
      saved_estimator_path = os.path.join(tmp_path, "saved_model")
      os.makedirs(saved_estimator_path)
      saved_estimator_path = estimator.export_saved_model(
          saved_estimator_path, receiver_fn
      ).decode("utf-8")
      return SavedModelInfo(
          path=saved_estimator_path,
          meta_graph_tags=["serve"],
          signature_def_key="predict",
          inference_df=pd.DataFrame(training_features),
          expected_results_df=estimator_preds_df,
          raw_results=None,
          raw_df=None,
      )
  
file_path: mlflow/tests/tensorflow/save_tf_estimator_model.py
imports: import argparse
  import collections
  import copy
  import os
  import pickle
  import iris_data_utils
  import numpy
  import pandas
  import pandas.testing
  import tensorflow
  from tensorflow import estimator
  import mlflow
  from mlflow.utils.file_utils import TempDir
label_desc: Use image node_id 2 for calling a global function with example usage: save_tf_categorical_model(tmp_path) and returns: SavedModelInfo