Uno de los enfoques no triviales más simples para la generación de texto es hacer lo que se implementa en el tutorial en la línea resaltada ###:

@tf.function
  def generate_one_step(self, inputs, states=None):
    input_chars = tf.strings.unicode_split(inputs, 'UTF-8')
    input_ids = self.ids_from_chars(input_chars).to_tensor()
    predicted_logits, states = self.model(inputs=input_ids, states=states,
                                          return_state=True)
    predicted_logits = predicted_logits[:, -1, :]
    predicted_logits = predicted_logits/self.temperature
    predicted_logits = predicted_logits + self.prediction_mask

    ### THE FOLLOWING LINE IS IMPORTANT ###
    predicted_ids = tf.random.categorical(predicted_logits, num_samples=1)

    predicted_ids = tf.squeeze(predicted_ids, axis=-1)
    predicted_chars = self.chars_from_ids(predicted_ids)
    return predicted_chars, states

Probablemente recibió una alerta porque esperaba que la generación fuera determinista, lo cual es cierto para el enfoque ingenuo: siempre devuelva el carácter/palabra/token/etc. más probable, pero no es así como funciona el lenguaje en absoluto. Si está interesado en los detalles, vaya a ver el curso de PNL de Stanford (se puede ver libremente en YouTube), de lo contrario, bueno, ahí lo tiene, hay aleatoriedad en el algoritmo.