def max_sequence_length(self, dataset_split):
    """Determine the maximum sequence length given a dataset_split.

    Args:
      dataset_split: A problem.DatasetSplit.

    Returns:
      The maximum length that a sequence can be for this dataset_split.
    """
    return {
        problem.DatasetSplit.TRAIN: 64,
        problem.DatasetSplit.EVAL: 128,
        problem.DatasetSplit.TEST: 128
    }[dataset_split]
Determine the maximum sequence length given a dataset_split.

    Args:
      dataset_split: A problem.DatasetSplit.

    Returns:
      The maximum length that a sequence can be for this dataset_split.