def METHOD_NAME(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