When running tensorflow-federated v.0.20.0 I get the error "tensorflow-federated.python.learning.models" has no attribute ModelWeights - nlp

Question

I tried running the tensorflow-federated tutorial on "Federated Learning for Text Generation" but keep getting AttributeErrors on tff.learning.models. I wonder if this is because I have an outdated version of tensorflow-federated 0.20.0. The following code is what I tried running:
NUM_ROUNDS = 5
# The state of the FL server, containing the model and optimization state.
state = fed_avg.initialize()
# Load our pre-trained Keras model weights into the global model state.
pre_trained_weights = tff.learning.models.ModelWeights(
    trainable=[v.numpy() for v in keras_model.trainable_weights],
    non_trainable=[v.numpy() for v in keras_model.non_trainable_weights]
)
state = fed_avg.set_model_weights(state, pre_trained_weights)
def keras_evaluate(state, round_num):
  # Take our global model weights and push them back into a Keras model to
  # use its standard `.evaluate()` method.
  keras_model = load_model(batch_size=BATCH_SIZE)
  keras_model.compile(
      loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
      metrics=[FlattenedCategoricalAccuracy()])
  model_weights = fed_avg.get_model_weights(state)
  model_weights.assign_weights_to(keras_model)
  loss, accuracy = keras_model.evaluate(example_dataset, steps=2, verbose=0)
  print('\tEval: loss={l:.3f}, accuracy={a:.3f}'.format(l=loss, a=accuracy))
for round_num in range(NUM_ROUNDS):
  print('Round {r}'.format(r=round_num))
  keras_evaluate(state, round_num)
  result = fed_avg.next(state, train_datasets)
  state = result.state
  train_metrics = result.metrics['client_work']['train']
  print('\tTrain: loss={l:.3f}, accuracy={a:.3f}'.format(
      l=train_metrics['loss'], a=train_metrics['accuracy']))
print('Final evaluation')
keras_evaluate(state, NUM_ROUNDS + 1)
Also the reason I am running v.0.20.0 is because this is the only version I got to work in my conda environment on my Linux system. Please come with comments on what I have done wrong I am very new to tensorflow-federated. Also since this is tutorial form tensorflow I am pretty sure the code should run if done correctly.
I tried to find documentation on tensorflow-federated.learning.models and there it says there should be a module called ModelWeights.