input: Write a function to generate a Machine Learning model in Python that can recognize handwritten digits from 0 to 9 using TensorFlow and Keras. Assume the input is already preprocessed as a flattened array of pixel values for the MNIST dataset. The function should return the compiled model.
  assert model.predict(X_train[:1]).argmax() == y_train[0].argmax()
  assert model.predict(X_test[:1]).argmax() == y_test[0].argmax()
output: def generate_model():
      import tensorflow as tf
      from tensorflow.keras import layers, models
      
      model = models.Sequential([
          layers.Dense(128, activation='relu', input_shape=(784,)),
          layers.Dense(64, activation='relu'),
          layers.Dense(10, activation='softmax')
      ])
      
      model.compile(optimizer='adam',
                    loss='sparse_categorical_crossentropy',
                    metrics=['accuracy'])
      
      return model