Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Add support for M2M100EncoderModel (aka NLLB, aka SONAR text encoder) #2541

Open
wants to merge 1 commit into
base: master
Choose a base branch
from
Open

Add support for M2M100EncoderModel (aka NLLB, aka SONAR text encoder) #2541

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 12 additions & 1 deletion sentence_transformers/models/Transformer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
from torch import nn
from transformers import AutoModel, AutoTokenizer, AutoConfig, T5Config, MT5Config
from transformers import AutoModel, AutoTokenizer, AutoConfig, T5Config, MT5Config, M2M100Config
import json
from typing import List, Dict, Optional, Union, Tuple
import os
Expand Down Expand Up @@ -61,6 +61,8 @@ def _load_model(self, model_name_or_path, config, cache_dir, **model_args):
self._load_t5_model(model_name_or_path, config, cache_dir, **model_args)
elif isinstance(config, MT5Config):
self._load_mt5_model(model_name_or_path, config, cache_dir, **model_args)
elif isinstance(config, M2M100Config):
self._load_m2m100_model(model_name_or_path, config, cache_dir, **model_args)
else:
self.auto_model = AutoModel.from_pretrained(
model_name_or_path, config=config, cache_dir=cache_dir, **model_args
Expand All @@ -75,6 +77,15 @@ def _load_t5_model(self, model_name_or_path, config, cache_dir, **model_args):
model_name_or_path, config=config, cache_dir=cache_dir, **model_args
)

def _load_m2m100_model(self, model_name_or_path, config, cache_dir, **model_args):
"""Loads the encoder model from M2M100 (aka NLLB, aka SONAR text encoder)"""
from transformers import M2M100EncoderModel

M2M100EncoderModel._keys_to_ignore_on_load_unexpected = ["decoder.*"]
self.auto_model = M2M100EncoderModel.from_pretrained(
model_name_or_path, config=config, cache_dir=cache_dir, **model_args
)

def _load_mt5_model(self, model_name_or_path, config, cache_dir, **model_args):
"""Loads the encoder model from T5"""
from transformers import MT5EncoderModel
Expand Down