Example how to use GeneratorBasedBuilder to train model Keras style

What I need help with / What I was wondering
I've created a Dataset by using tfds.core.GeneratorBasedBuilder and want to train a model with it in Keras style. I followed the tutorial: https://github.com/tensorflow/datasets/blob/master/docs/add_dataset.md#datasetbuilder

I was wondering how to use this dataset to train a model Keras style with TF 2.0.
I want to make the model work with minimal code that looks like this:

model = tf.keras.models.Sequential([
  tf.keras.layers.Flatten(input_shape=(28, 28)),
  tf.keras.layers.Dense(128, activation='relu'),
  tf.keras.layers.Dropout(0.2),
  tf.keras.layers.Dense(10, activation='softmax')
])

model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])

model.fit_generator(ds_train, epochs=5)

With Keras, I was used to using model.fit_generator() for this purpose, however that will result (not unexpectedly) in the error:

ValueError: Output of generator should be a tuple `(x, y, sample_weight)` or `(x, y)`. Found: {'audio_description': <tf.Tensor: id=83383, shape=(5,), dtype=string, numpy=array([b'61_31', b'34_2', b'107_2', b'212_31', b'14_2'], dtype=object)>, 'audio': <tf.Tensor: id=83382, shape=(5, 2), dtype=float32, numpy=
array([[-0.3,  0.2],
       [ 0.3,  0.2],
       [ 0.3,  0.2],
       [-0.3,  0.2],
       [ 0.3,  0.2]], dtype=float32)>, 'label': <tf.Tensor: id=83384, shape=(5,), dtype=int64, numpy=array([1, 0, 0, 1, 0])>}

Epoch 1/5

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-52-f2489416f405> in <module>
----> 1 model.fit_generator(ds_train, epochs=5)

/usr/local/lib/python3.5/dist-packages/tensorflow/python/keras/engine/training.py in fit_generator(self, generator, steps_per_epoch, epochs, verbose, callbacks, validation_data, validation_steps, validation_freq, class_weight, max_queue_size, workers, use_multiprocessing, shuffle, initial_epoch)
   1513         shuffle=shuffle,
   1514         initial_epoch=initial_epoch,
-> 1515         steps_name='steps_per_epoch')
   1516 
   1517   def evaluate_generator(self,

/usr/local/lib/python3.5/dist-packages/tensorflow/python/keras/engine/training_generator.py in model_iteration(model, data, steps_per_epoch, epochs, verbose, callbacks, validation_data, validation_steps, validation_freq, class_weight, max_queue_size, workers, use_multiprocessing, shuffle, initial_epoch, mode, batch_size, steps_name, **kwargs)
    211     step = 0
    212     while step < target_steps:
--> 213       batch_data = _get_next_batch(generator, mode)
    214       if batch_data is None:
    215         if is_dataset:

/usr/local/lib/python3.5/dist-packages/tensorflow/python/keras/engine/training_generator.py in _get_next_batch(generator, mode)
    363       raise ValueError('Output of generator should be '
    364                        'a tuple `(x, y, sample_weight)` '
--> 365                        'or `(x, y)`. Found: ' + str(generator_output))
    366 
    367   if len(generator_output) < 1 or len(generator_output) > 3:

ValueError: Output of generator should be a tuple `(x, y, sample_weight)` or `(x, y)`. Found: {'audio_description': <tf.Tensor: id=83383, shape=(5,), dtype=string, numpy=array([b'61_31', b'34_2', b'107_2', b'212_31', b'14_2'], dtype=object)>, 'audio': <tf.Tensor: id=83382, shape=(5, 2), dtype=float32, numpy=
array([[-0.3,  0.2],
       [ 0.3,  0.2],
       [ 0.3,  0.2],
       [-0.3,  0.2],
       [ 0.3,  0.2]], dtype=float32)>, 'label': <tf.Tensor: id=83384, shape=(5,), dtype=int64, numpy=array([1, 0, 0, 1, 0])>}

Question: How to use a DatasetBuilder with model.fit_generator in a single/few lines of code Keras style?

What I've tried so far
Following the expert introduction to TF 2.0, I got this to work with minimal changes:

from tensorflow.keras.layers import Dense, Flatten, Conv2D
from tensorflow.keras import Model


class MyModel(Model):
    def __init__(self):
        super(MyModel, self).__init__()
        #self.conv1 = Conv2D(32, 3, activation='relu')
        #self.flatten = Flatten()
        self.d1 = Dense(2, activation='relu')
        self.d2 = Dense(2, activation='softmax')

    def call(self, x):
        #x = self.conv1(x)
        #x = self.flatten(x)
        x = self.d1(x)
        return self.d2(x)

model = MyModel()

loss_object = tf.keras.losses.SparseCategoricalCrossentropy()
optimizer = tf.keras.optimizers.Adam()

train_loss = tf.keras.metrics.Mean(name='train_loss')
train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='train_accuracy')

test_loss = tf.keras.metrics.Mean(name='test_loss')
test_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='test_accuracy')

@tf.function
def train_step(images, labels):
    with tf.GradientTape() as tape:
        predictions = model(images)
        loss = loss_object(labels, predictions)
    gradients = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))

    train_loss(loss)
    train_accuracy(labels, predictions)
    
@tf.function
def test_step(images, labels):
    predictions = model(images)
    t_loss = loss_object(labels, predictions)

    test_loss(t_loss)
    test_accuracy(labels, predictions)

EPOCHS = 200

for epoch in range(EPOCHS):
    for features in ds_train:
        train_step(features["audio"], features["label"])
        
    for features in ds_train:
        test_step(features["audio"], features["label"])

    template = 'Epoch {}, Loss: {}, Accuracy: {}, Test Loss: {}, Test Accuracy: {}'
    print (template.format(epoch+1,
                                                 train_loss.result(),
                                                 train_accuracy.result()*100,
                                                 test_loss.result(),
                                                 test_accuracy.result()*100))

However, that the destroys the new philosophy of easy to use TF 2.0 with Keras.

It would be nice if...
Please provide an example, Keras style, on how to easily use DatasetBuilder / GeneratorBasedBuilder to train a model and not stopping at for features in ds_train:.

Environment information
(if applicable)

• Operating System: Ubuntu 18.04 running Docker container: docker run -it --runtime=nvidia --rm -v /home/notebooks:/tf/notebooks -p 8889:8888 tensorflow/tensorflow:2.0.0a0-gpu-py3-jupyter
• Python version (in Docker): 3.5
• tensorflow-datasets version: '1.0.2' (tensorflow_datasets.version.__version__)
• tensorflow-gpu version: 2.0.0-alpha

Start with docs/add_dataset.md#datasetbuilder and the issue's ds_train and model.fit_generator examples. Review how the audio and label entries are exposed by the feature dictionary, then define done as a concise TF 2.0 Keras-style training example that addresses the reported generator error.