input: Write a function to construct a neural network model in Python using Keras to correctly classify the MNIST data set. The function should create a model with an input layer of shape (784,), a hidden layer with 128 units and 'relu' activation, and an output layer with 10 units and 'softmax' activation. Finally, the model should be compiled using the 'adam' optimizer and 'sparse_categorical_crossentropy' loss function.
  
  def mnist_model():
      import keras
  
      # Input layer
      input_layer = keras.layers.Input(shape=(784,))
  
      # Hidden layer
      hidden_layer = keras.layers.Dense(128, activation='relu')(input_layer)
  
      # Output layer
      output_layer = keras.layers.Dense(10, activation='softmax')(hidden_layer)
  
      # Model
      model = keras.Model(inputs=input_layer, outputs=output_layer)
  
      # Compile model
      model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])
      
      return model
  
  assert type(mnist_model()) == keras.models.Model
  assert mnist_model().output_shape == (None, 10)
  assert mnist_model().layers[0].input_shape == [(None, 784)]
  assert mnist_model().layers[1].units == 128
  assert mnist_model().layers[1].activation == keras.activations.relu
  assert mnist_model().layers[2].units == 10
  assert mnist_model().layers[2].activation == keras.activations.softmax
  assert mnist_model().loss == 'sparse_categorical_crossentropy'
  assert mnist_model().optimizer.__class__ == keras.optimizers.Adam
output: def mnist_model():
      import keras
  
      input_layer = keras.layers.Input(shape=(784,))
      hidden_layer = keras.layers.Dense(128, activation='relu')(input_layer)
      output_layer = keras.layers.Dense(10, activation='softmax')(hidden_layer)
      model = keras.Model(inputs=input_layer, outputs=output_layer)
      model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])
      
      return model