input: Write a function to generate a machine learning model in Python that can recognize handwritten digits from 0 to 9 using the MNIST dataset. The function should load the dataset, preprocess the data, build a neural network model with specified layers, compile the model, and fit the model to the training data. The function should return the trained model.
  assert model.evaluate(X_test, y_test)[1] > 0.90
  assert model.fit(X_train, y_train, validation_data=(X_test, y_test), epochs=5, batch_size=128) is not None
  assert model.predict(X_test[0:1]).shape == (1, 10)
  assert model.predict(X_test[0:1]).dtype == 'float32'
  assert model.compile(loss='categorical_crossentropy', optimizer='adam', metrics=['accuracy']) is None
output: def generate_model():
      from tensorflow.keras.datasets import mnist
      from tensorflow.keras.models import Sequential
      from tensorflow.keras.layers import Dense, Flatten
      from tensorflow.keras.utils import to_categorical
  
      (X_train, y_train), (X_test, y_test) = mnist.load_data()
      X_train = X_train.reshape(-1, 28*28) / 255.0
      X_test = X_test.reshape(-1, 28*28) / 255.0
      y_train = to_categorical(y_train, 10)
      y_test = to_categorical(y_test, 10)
  
      model = Sequential([
          Flatten(input_shape=(28*28,)),
          Dense(128, activation='relu'),
          Dense(64, activation='relu'),
          Dense(10, activation='softmax')
      ])
  
      model.compile(loss='categorical_crossentropy', optimizer='adam', metrics=['accuracy'])
      model.fit(X_train, y_train, validation_data=(X_test, y_test), epochs=5, batch_size=128)
  
      return model