def extend_position_embedding(model, max_position):
    """This function extends the position embedding weights of a model loaded from a checkpoint.
    It assumes the new max position is bigger than the original max length.

    Arguments:
        model: required: a transformer model
        max_position: required: an integer determining new position embedding size
    Return:
        model: updated model; in which position embedding weights have been extended based on new size
    """

    if hasattr(model, 'bert'):
        original_max_position = model.bert.embeddings.position_embeddings.weight.size(0)
        assert max_position > original_max_position
        extend_multiples = max(1, max_position // original_max_position)
        model.bert.embeddings.position_embeddings.weight.data = model.bert.embeddings.position_embeddings.weight.repeat(
            extend_multiples,
            1)
    elif hasattr(model, 'roberta'):
        # RoBERTa has positions 0 & 1 reserved, so embedding size is max position + 2
        original_max_position, embed_size = model.roberta.embeddings.position_embeddings.weight.shape
        original_max_position -= 2
        extend_multiples = max(1, max_position // original_max_position)
        assert max_position > original_max_position
        max_position += 2
        extended_position_embedding = model.roberta.embeddings.position_embeddings.weight.new_empty(
            max_position,
            embed_size)
        k = 2
        for i in range(extend_multiples):
            extended_position_embedding[k:(
                k + original_max_position
            )] = model.roberta.embeddings.position_embeddings.weight[2:]
            k += original_max_position
        model.roberta.embeddings.position_embeddings.weight.data = extended_position_embedding
    else:
        raise ValueError(
            'Please extend \"extend_position_embedding\" function to support your model type. It currently only supports \"bert\" & \"roberta\"!'
        )

    model.config.max_position_embeddings = max_position
    print(f'Extended position embeddings to {original_max_position * extend_multiples}')

    return model