File size: 932 Bytes
ded35c4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 |
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)
|