3008753be172bc0bae88aa14b00ce860233c3416,official/nlp/tasks/utils.py,,get_encoder_from_hub,#,24

Before Change

  hub_layer = hub.KerasLayer(hub_module, trainable=True)
  pooled_output, sequence_output = hub_layer(
      [input_word_ids, input_mask, input_type_ids])
  return tf.keras.Model(
      inputs=[input_word_ids, input_mask, input_type_ids],
      outputs=[sequence_output, pooled_output])


def predict(predict_step_fn: Callable[[Any], Any],
            aggregate_fn: Callable[[Any, Any], Any], dataset: tf.data.Dataset):

After Change

    logging.info("Use the new hub module with dict as input/output.")
    output_dict = hub_layer(dict_input)

  return tf.keras.Model(inputs=dict_input, outputs=output_dict)


def predict(predict_step_fn: Callable[[Any], Any],
            aggregate_fn: Callable[[Any, Any], Any], dataset: tf.data.Dataset):

Instances