code: def build_model(
      input_shape: List[int], num_classes: int
  ) -> tf.keras.Model:
      """Builds a simple Keras model for test."""
      inputs = tf.keras.layers.Input(shape=input_shape)
      if len(input_shape) == 3:  # Image inputs.
          outputs = tf.keras.layers.GlobalAveragePooling2D()(inputs)
          outputs = tf.keras.layers.Dense(
              num_classes, activation="softmax"
          )(outputs)
      elif len(input_shape) == 1:  # Text inputs.
          outputs = tf.keras.layers.Dense(
              num_classes, activation="softmax"
          )(inputs)
      else:
          raise ValueError(
              "Model inputs should be 2D tensor or 4D tensor."
          )
  
      model = tf.keras.Model(inputs=inputs, outputs=outputs)
      return model
  
file_path: mediapipe/mediapipe/model_maker/python/core/utils/test_util.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from typing import Sequence
  from typing import Dict, List, Union
  import numpy
  import tensorflow
  from mediapipe.model_maker.python.core.data import dataset
  from mediapipe.model_maker.python.core.utils import model_util
label_desc: Use image node_id 3 for calling a global function with example usage: build_model(input_shape, num_classes) and returns: model