Skip to content

Update HF interface - fix from_config issue #2293

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

Merged
mravanelli merged 1 commit into from
Dec 11, 2023
Merged

Update HF interface - fix from_config issue #2293

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
67 changes: 32 additions & 35 deletions speechbrain/lobes/models/huggingface_transformers/huggingface.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,7 @@
AutoModel,
AutoModelWithLMHead,
AutoModelForSeq2SeqLM,
AutoModelForCausalLM,
)

logger = logging.getLogger(__name__)
Expand Down Expand Up @@ -65,14 +66,18 @@ class HFTransformersInterface(nn.Module):
save directory of the downloaded model.
for_pretraining: bool (default: False)
If True, build the model for pretraining
with_lm_head: bool (default: False)
with_lm_head : bool (default: False)
If True, build the model with lm_head
with_casual_lm : bool (default: False)
If True, build casual lm model
seq2seqlm : bool (default: False)
If True, build a sequence-to-sequence model with lm_head
quantization_config : dict (default: None)
Quantization config, extremely useful for deadling with LLM
freeze : bool (default: True)
If True, the model is frozen. If False, the model will be trained
alongside with the rest of the pipeline.
cache_dir: str or Path (default: None)
cache_dir : str or Path (default: None)
Location of HuggingFace cache for storing pre-trained models, to which symlinks are created.

Example
Expand All @@ -88,7 +93,9 @@ def __init__(
save_path="",
for_pretraining=False,
with_lm_head=False,
with_casual_lm=False,
seq2seqlm=False,
quantization_config=None,
freeze=False,
cache_dir="pretrained_models",
**kwarg,
Expand All @@ -101,26 +108,24 @@ def __init__(
)

self.config = self.override_config(self.config)
self.quantization_config = quantization_config

self.for_pretraining = for_pretraining
self.with_lm_head = with_lm_head
self.seq2seqlm = seq2seqlm
if for_pretraining:
model = AutoModelForPreTraining.from_config(self.config)

if self.for_pretraining:
self.auto_class = AutoModelForPreTraining
elif with_lm_head:
model = AutoModelWithLMHead.from_config(self.config)
self.auto_class = AutoModelWithLMHead
elif with_casual_lm:
self.auto_class = AutoModelForCausalLM
elif seq2seqlm:
model = AutoModelForSeq2SeqLM.from_config(self.config)
self.auto_class = AutoModelForSeq2SeqLM
else:
model = AutoModel.from_config(self.config)
self.auto_class = AutoModel

# Download model
self._from_pretrained(
source,
config=self.config,
model=model,
save_path=save_path,
cache_dir=cache_dir,
source, save_path=save_path, cache_dir=cache_dir,
)

# Prepare for training, fine-tuning, or inference
Expand All @@ -135,7 +140,7 @@ def __init__(
self.model.train()

def _from_pretrained(
self, source, config, model, save_path, cache_dir,
self, source, save_path, cache_dir,
):
"""This function manages the source checking and loading of the params.

Expand All @@ -147,39 +152,31 @@ def _from_pretrained(
---------
source : str
HuggingFace hub name: e.g "facebook/wav2vec2-large-lv60"
config : AutoConfig
HuggingFace generic configuration class.
model: AutoModel
HuggingFace generic model class.
save_path : str
Path (dir) of the downloaded model.
cache_dir : str
Path (dir) in which a downloaded pretrained model configuration should be cached.
"""
is_sb, ckpt_file, is_local = self._check_model_source(source, save_path)

if is_sb or self.for_pretraining:
self.model = self.auto_class.from_config(self.config)

if is_sb:
config = config.from_pretrained(source, cache_dir=save_path)
self.model = model(config)
self.model.gradient_checkpointing_disable() # Required by DDP
# fetch the checkpoint file
ckpt_full_path = fetch(
filename=ckpt_file,
source=source,
savedir=save_path,
cache_dir=cache_dir,
filename=ckpt_file, source=source, savedir=save_path,
)
# We transfer the parameters from the checkpoint.
self._load_sb_pretrained_parameters(ckpt_full_path)
else:
if self.for_pretraining:
# For now, we don't need to load pretrained model for pretraining
# To be modified in the future to support more complicated scenerios
# For example fine-tuning in the SSL manner
self.model = model
else:
self.model = model.from_pretrained(
source, config=config, cache_dir=save_path
)
elif not self.for_pretraining:
self.model = self.auto_class.from_pretrained(
source,
config=self.config,
cache_dir=save_path,
quantization_config=self.quantization_config,
)

def _check_model_source(self, path, save_path):
"""Checks if the pretrained model has been trained with SpeechBrain and
Expand Down