|
from pylate import models |
|
|
|
|
|
|
|
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) |
|
|
