Skip to content

cross platform from_pretrained #20538

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
ArthurZucker merged 10 commits into from
Dec 5, 2022
Merged

cross platform from_pretrained #20538

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
60c0ee6
add support for `from_pt`
ArthurZucker Dec 1, 2022
cf9fbb3
add tf_flax utility file
ArthurZucker Dec 1, 2022
e924d1d
Update src/transformers/modeling_tf_flax_utils.py
ArthurZucker Dec 2, 2022
960e145
remove flax related modifications
ArthurZucker Dec 2, 2022
437db80
add test
ArthurZucker Dec 2, 2022
b7b645a
Merge branch 'tf-from-sharded-pt' of https://github.com/ArthurZucker/…
ArthurZucker Dec 2, 2022
22cf129
remove FLAX related commits
ArthurZucker Dec 2, 2022
fa1d480
fixup
ArthurZucker Dec 2, 2022
2bf1f31
remove safetensor todos
ArthurZucker Dec 2, 2022
a43dbc0
revert deletion
ArthurZucker Dec 5, 2022
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
20 changes: 20 additions & 0 deletions src/transformers/modeling_tf_flax_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,20 @@
# coding=utf-8
# Copyright 2022 The HuggingFace Inc. team.
# Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved.
Comment thread
ArthurZucker marked this conversation as resolved.
Outdated
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
""" Flax - TF 2.0 general utilities."""


def load_flax_checkpoint_in_tf2_model(model, resolved_archive_file):
return NotImplementedError
79 changes: 67 additions & 12 deletions src/transformers/modeling_tf_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,8 @@
from .tf_utils import shape_list
from .utils import (
DUMMY_INPUTS,
FLAX_WEIGHTS_INDEX_NAME,
FLAX_WEIGHTS_NAME,
SAFE_WEIGHTS_INDEX_NAME,
SAFE_WEIGHTS_NAME,
TF2_WEIGHTS_INDEX_NAME,
Expand Down Expand Up @@ -2395,6 +2397,9 @@ def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
from_pt: (`bool`, *optional*, defaults to `False`):
Load the model weights from a PyTorch state_dict save file (see docstring of
`pretrained_model_name_or_path` argument).
from_flax: (`bool`, *optional*, defaults to `False`):
Comment thread
ArthurZucker marked this conversation as resolved.
Outdated
Load the model weights from a Flax state_dict save file (see docstring of
`pretrained_model_name_or_path` argument).
ignore_mismatched_sizes (`bool`, *optional*, defaults to `False`):
Whether or not to raise an error if some of the weights from the checkpoint do not have the same size
as the weights of the model (if for instance, you are instantiating a model with 10 labels from a
Expand Down Expand Up @@ -2470,6 +2475,7 @@ def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
config = kwargs.pop("config", None)
cache_dir = kwargs.pop("cache_dir", None)
from_pt = kwargs.pop("from_pt", False)
from_flax = kwargs.pop("from_flax", False)
ignore_mismatched_sizes = kwargs.pop("ignore_mismatched_sizes", False)
force_download = kwargs.pop("force_download", False)
resume_download = kwargs.pop("resume_download", False)
Expand Down Expand Up @@ -2531,14 +2537,23 @@ def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
if pretrained_model_name_or_path is not None:
pretrained_model_name_or_path = str(pretrained_model_name_or_path)
is_local = os.path.isdir(pretrained_model_name_or_path)
if os.path.isdir(pretrained_model_name_or_path):
if is_local:
if from_pt and os.path.isfile(os.path.join(pretrained_model_name_or_path, WEIGHTS_NAME)):
# Load from a PyTorch checkpoint in priority if from_pt
archive_file = os.path.join(pretrained_model_name_or_path, WEIGHTS_NAME)
elif from_pt and os.path.isfile(os.path.join(pretrained_model_name_or_path, WEIGHTS_INDEX_NAME)):
# Load from a sharded PyTorch checkpoint
archive_file = os.path.join(pretrained_model_name_or_path, WEIGHTS_INDEX_NAME)
is_sharded = True
elif from_flax and os.path.isfile(os.path.join(pretrained_model_name_or_path, FLAX_WEIGHTS_NAME)):
# Load from a Flax checkpoint
archive_file = os.path.join(pretrained_model_name_or_path, FLAX_WEIGHTS_NAME)
elif from_flax and os.path.isfile(
os.path.join(pretrained_model_name_or_path, FLAX_WEIGHTS_INDEX_NAME)
):
# Load from a sharded Flax checkpoint
archive_file = os.path.join(pretrained_model_name_or_path, FLAX_WEIGHTS_INDEX_NAME)
is_sharded = True
elif is_safetensors_available() and os.path.isfile(
os.path.join(pretrained_model_name_or_path, SAFE_WEIGHTS_NAME)
):
Expand All @@ -2559,7 +2574,17 @@ def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
archive_file = os.path.join(pretrained_model_name_or_path, TF2_WEIGHTS_INDEX_NAME)
is_sharded = True
# At this stage we don't have a weight file so we will raise an error.
elif os.path.join(pretrained_model_name_or_path, WEIGHTS_NAME):
elif os.path.isfile(
os.path.join(pretrained_model_name_or_path, subfolder, FLAX_WEIGHTS_NAME)
) or os.path.isfile(os.path.join(pretrained_model_name_or_path, FLAX_WEIGHTS_INDEX_NAME)):
raise EnvironmentError(
f"Error no file named {TF2_WEIGHTS_NAME} found in directory"
f" {pretrained_model_name_or_path} but there is a file for Flax weights. Use `from_flax=True`"
" to load this model from those weights."
)
elif os.path.isfile(os.path.join(pretrained_model_name_or_path, WEIGHTS_NAME)) or os.path.isfile(
os.path.join(pretrained_model_name_or_path, WEIGHTS_INDEX_NAME)
):
raise EnvironmentError(
f"Error no file named {TF2_WEIGHTS_NAME} found in directory {pretrained_model_name_or_path} "
"but there is a file for PyTorch weights. Use `from_pt=True` to load this model from those "
Expand All @@ -2570,19 +2595,15 @@ def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
f"Error no file named {TF2_WEIGHTS_NAME} or {WEIGHTS_NAME} found in directory "
f"{pretrained_model_name_or_path}."
)
elif os.path.isfile(pretrained_model_name_or_path):
archive_file = pretrained_model_name_or_path
is_local = True
elif os.path.isfile(pretrained_model_name_or_path + ".index"):
archive_file = pretrained_model_name_or_path + ".index"
is_local = True
Comment on lines -2573 to -2578

@sgugger sgugger Dec 5, 2022

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This code shouldn't be removed, to preserve compatibility with PreTrainedModel.from_pretrained(path_to_a_model_path, config = config)

elif is_remote_url(pretrained_model_name_or_path):
filename = pretrained_model_name_or_path
resolved_archive_file = download_url(pretrained_model_name_or_path)
else:
# set correct filename
if from_pt:
filename = WEIGHTS_NAME
elif from_flax:
filename = FLAX_WEIGHTS_NAME
elif is_safetensors_available():
filename = SAFE_WEIGHTS_NAME
else:
Expand Down Expand Up @@ -2612,7 +2633,7 @@ def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
resolved_archive_file = cached_file(
pretrained_model_name_or_path, SAFE_WEIGHTS_INDEX_NAME, **cached_file_kwargs
)
if resolved_archive_file is not None:
if resolved_archive_file is not None: # TODO Arthur, add safetensors support for sharded tf
is_sharded = True
raise NotImplementedError(
"Support for sharded checkpoints using safetensors is coming soon!"
Expand All @@ -2630,6 +2651,20 @@ def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
)
if resolved_archive_file is not None:
is_sharded = True
if resolved_archive_file is None and filename == WEIGHTS_NAME:
# Maybe the checkpoint is sharded, we try to grab the index name in this case.
resolved_archive_file = cached_file(
pretrained_model_name_or_path, WEIGHTS_INDEX_NAME, **cached_file_kwargs
)
if resolved_archive_file is not None:
is_sharded = True
Comment thread
ArthurZucker marked this conversation as resolved.
if resolved_archive_file is None and filename == FLAX_WEIGHTS_NAME:
# Maybe the checkpoint is sharded, we try to grab the index name in this case.
resolved_archive_file = cached_file(
pretrained_model_name_or_path, FLAX_WEIGHTS_INDEX_NAME, **cached_file_kwargs
)
if resolved_archive_file is not None:
is_sharded = True
if resolved_archive_file is None:
# Otherwise, maybe there is a PyTorch or Flax model file. We try those to give a helpful error
# message.
Expand All @@ -2644,10 +2679,16 @@ def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
f" {TF2_WEIGHTS_NAME} but there is a file for PyTorch weights. Use `from_pt=True` to"
" load this model from those weights."
)
else:
elif has_file(pretrained_model_name_or_path, FLAX_WEIGHTS_NAME, **has_file_kwargs):
raise EnvironmentError(
f"{pretrained_model_name_or_path} does not appear to have a file named"
f" {TF2_WEIGHTS_NAME} or {WEIGHTS_NAME}."
f" {WEIGHTS_NAME} but there is a file for TensorFlow weights. Use `from_tf=True` to"
" load this model from those weights."
)
else:
raise EnvironmentError(
f"{pretrained_model_name_or_path} does not appear to have a file named {WEIGHTS_NAME},"
f" {TF2_WEIGHTS_NAME}, or {FLAX_WEIGHTS_NAME}."
)

except EnvironmentError:
Expand All @@ -2661,7 +2702,8 @@ def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
f"Can't load the model for '{pretrained_model_name_or_path}'. If you were trying to load it"
" from 'https://huggingface.co/models', make sure you don't have a local directory with the"
f" same name. Otherwise, make sure '{pretrained_model_name_or_path}' is the correct path to a"
f" directory containing a file named {TF2_WEIGHTS_NAME} or {WEIGHTS_NAME}."
f" directory containing a file named {WEIGHTS_NAME}, {TF2_WEIGHTS_NAME}, {TF_WEIGHTS_NAME} or"
f" {FLAX_WEIGHTS_NAME}."
)
if is_local:
logger.info(f"loading weights file {archive_file}")
Expand Down Expand Up @@ -2689,6 +2731,7 @@ def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
_commit_hash=commit_hash,
)

# TODO this might not support the sharded safetensors
safetensors_from_pt = False
if filename == SAFE_WEIGHTS_NAME:
with safe_open(resolved_archive_file, framework="tf") as f:
Expand Down Expand Up @@ -2717,6 +2760,18 @@ def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
return load_pytorch_checkpoint_in_tf2_model(
model, resolved_archive_file, allow_missing_keys=True, output_loading_info=output_loading_info
)
elif from_flax:
try:
from .modeling_tf_flax_utils import load_flax_checkpoint_in_tf2_model

model = load_flax_checkpoint_in_tf2_model(model, resolved_archive_file)
except ImportError:
logger.error(
"Loading a Flax model in PyTorch, requires both PyTorch and Flax to be installed. Please see"
" https://pytorch.org/ and https://flax.readthedocs.io/en/latest/installation.html for"
" installation instructions."
)
raise
elif safetensors_from_pt:
from .modeling_tf_pytorch_utils import load_pytorch_state_dict_in_tf2_model

Expand Down