t5x.eval binary#

Runs training- and inference-evaluation on a T5X-compatible model.

class t5x.eval.InferenceEvaluator(infer_eval_dataset_cfg, inference_evaluator_cls, model, partitioner, log_dir=None, verify_matching_vocabs_fn=<function verify_matching_vocabs>)[source]#

Runs evaluation of the model against a given SeqIo task.

evaluate(train_state, train_state_axes)[source]#

Runs the prediction based inference eval.

Parameters:

  • train_state – Training state to run evaluation of.

  • train_state_axes – partitioning info for the train state to be used.

Returns:

A dictionary of training eval metrics.

class t5x.eval.SummarizeConfigFn(*args, **kwargs)[source]#
t5x.eval.evaluate(*, model, dataset_cfg, restore_checkpoint_cfg, partitioner, output_dir, inference_evaluator_cls=<class 'seqio.evaluation.Evaluator'>, training_evaluator_cls=None, summarize_config_fn=<function summarize_gin_config>, train_state_initializer_cls=<class 't5x.utils.TrainStateInitializer'>, train_eval_get_dataset_fn=<function get_training_eval_datasets>, fallback_init_rng=None, use_orbax=False)[source]#

Evaluation function.

Parameters:

  • model – The model object to use for inference.

  • dataset_cfg – Specification for the dataset to infer based on.

  • restore_checkpoint_cfg – Specification for the model parameter checkpoint to load.

  • partitioner – Partitioner for the model parameters and data across devices.

  • output_dir – Path to directory to write temporary files and final results.

  • inference_evaluator_cls – seqio.Evaluator class to use for inference evaluation, potentially with bound configuration args.

  • training_evaluator_cls – an optional Trainer class to use for training evaluation, potentially with bound configuration args.

  • summarize_config_fn – A function that takes in the model directory, an optional SummaryWriter, and the step number, and writes a summary of the configuration. SummaryWriter will be None in most cases.

  • train_state_initializer_cls – t5x.utils.TrainStateInitializer class for initializing partitioned TrainState from checkpoints or scratch.

  • train_eval_get_dataset_fn – Optional callable use to get the train-eval datasets based on the DatasetConfig and shard information. If missing, it defaults to utils.get_training_eval_datasets.

  • fallback_init_rng – A random seed used for parameter initialization during model re-loading when utils.RestoreCheckpointConfig.fallback_to_scratch is set to True. If None, parameter initialization is not allowed during model loading and having fallback_to_scratch enabled will result in an error.

  • use_orbax – if True, uses Orbax for checkpointing. Experimental feature.