benchmarks/model_benchmark/bert_benchmark.py/load_model
def load_model():
    if FLAGS.model_size not in MODEL_SIZE_MAP.keys():
        raise KeyError(
            f"`model_size` must be one of {MODEL_SIZE_MAP.keys()}, but "
            f"received {FLAGS.model_size}."
        )
    return keras_nlp.models.BertClassifier.from_preset(
        MODEL_SIZE_MAP[FLAGS.model_size], num_classes=2
    )