input: def classify_mnist_model():
      """ Construct a neural network model in Python to classify the MNIST data set correctly.
      """
      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 constructed model
      return model
  
      # Example:
      # >>> mnist_model = classify_mnist_model()
      # >>> mnist_model.summary()
  
      # Expected Output:
      # Model: "model"
      # _________________________________________________________________
      # Layer (type)                 Output Shape              Param #   
      # =================================================================
      # input_1 (InputLayer)         [(None, 784)]             0         
      # _________________________________________________________________
      # dense (Dense)                (None, 128)               100480    
      # _________________________________________________________________
      # dense_1 (Dense)              (None, 10)                1290      
      # =================================================================
      # Total params: 101,770
      # Trainable params: 101,770
      # Non-trainable params: 0
      # _________________________________________________________________
  
      # Example:
      # >>> mnist_model.predict([[0]*784])
  
      # Expected Output:
      # array([[0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1]], dtype=float32)
output:     import numpy as np
      from keras.datasets import mnist
      from keras.utils import to_categorical
  
      # Load MNIST dataset
      (x_train, y_train), (x_test, y_test) = mnist.load_data()
  
      # Preprocess the data
      x_train = x_train.reshape(-1, 784).astype('float32') / 255
      x_test = x_test.reshape(-1, 784).astype('float32') / 255
      y_train = to_categorical(y_train, 10)
      y_test = to_categorical(y_test, 10)
  
      # Train the model
      model.fit(x_train, y_train, epochs=5, batch_size=32, validation_data=(x_test, y_test))
  
      # Return the trained model
      return model