input: Write a function to create a neural network model using TensorFlow and Python for classifying hand-written digits from the MNIST dataset. The function should include parameters for learning rate, training epochs, batch size, and display step. It should also print the training cost for each epoch and the final accuracy on the test set.
  assert create_mnist_model()["accuracy"] > 0.9
output: def create_mnist_model(learning_rate=0.001, training_epochs=10, batch_size=100, display_step=1):
      import tensorflow as tf
      from tensorflow.keras import layers, models
      from tensorflow.keras.datasets import mnist
      from tensorflow.keras.utils import to_categorical
  
      (x_train, y_train), (x_test, y_test) = mnist.load_data()
      x_train, x_test = x_train / 255.0, x_test / 255.0
      x_train = x_train[..., tf.newaxis]
      x_test = x_test[..., tf.newaxis]
      y_train = to_categorical(y_train, 10)
      y_test = to_categorical(y_test, 10)
  
      model = models.Sequential([
          layers.Conv2D(32, kernel_size=(3, 3), activation='relu', input_shape=(28, 28, 1)),
          layers.MaxPooling2D(pool_size=(2, 2)),
          layers.Flatten(),
          layers.Dense(128, activation='relu'),
          layers.Dense(10, activation='softmax')
      ])
  
      model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=learning_rate),
                    loss='categorical_crossentropy',
                    metrics=['accuracy'])
  
      model.fit(x_train, y_train, batch_size=batch_size, epochs=training_epochs, verbose=1, validation_data=(x_test, y_test))
  
      test_loss, test_accuracy = model.evaluate(x_test, y_test, verbose=0)
      print(f"Final Accuracy on Test Set: {test_accuracy:.4f}")
  
      return {"accuracy": test_accuracy}