from pylate import models
# pip install pylate==1.1.2


def get_pylate_model(language: str, train: bool) -> 'pylate.models.ColBERT':
    """
    Configures the antoinelouis/colbert-xm model for usage with PyLate
    See discussion here: https://github.com/lightonai/pylate/discussions/50#discussioncomment-10691630
    For language, use a code from https://huggingface.co/facebook/xmod-base#languages
    """
    colbert_model = models.ColBERT(model_name_or_path='antoinelouis/colbert-xm')
    backbone = colbert_model[0].auto_model

    if backbone.__class__.__name__.lower().startswith("xmod"):
        backbone.set_default_language(language)

        if train:
            backbone.freeze_embeddings_and_language_adapters()

    training = '' if not train else '_train'
    colbert_model.save_pretrained(f'pylate-colbert-xm-{language}{training}/')


if __name__ == '__main__':
    get_pylate_model(language='de_DE', train=False)