803 lines
36 KiB
Python
803 lines
36 KiB
Python
# coding=utf-8
|
|
# Copyright 2019 HuggingFace Inc.
|
|
#
|
|
# 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.
|
|
|
|
|
|
from __future__ import annotations
|
|
|
|
import inspect
|
|
import json
|
|
import os
|
|
import random
|
|
import tempfile
|
|
import unittest
|
|
import unittest.mock as mock
|
|
|
|
from huggingface_hub import HfFolder, Repository, delete_repo, snapshot_download
|
|
from requests.exceptions import HTTPError
|
|
|
|
from transformers import is_tf_available, is_torch_available
|
|
from transformers.configuration_utils import PretrainedConfig
|
|
from transformers.testing_utils import ( # noqa: F401
|
|
TOKEN,
|
|
USER,
|
|
CaptureLogger,
|
|
_tf_gpu_memory_limit,
|
|
is_pt_tf_cross_test,
|
|
is_staging_test,
|
|
require_safetensors,
|
|
require_tf,
|
|
require_torch,
|
|
slow,
|
|
)
|
|
from transformers.utils import (
|
|
SAFE_WEIGHTS_INDEX_NAME,
|
|
SAFE_WEIGHTS_NAME,
|
|
TF2_WEIGHTS_INDEX_NAME,
|
|
TF2_WEIGHTS_NAME,
|
|
logging,
|
|
)
|
|
|
|
|
|
logger = logging.get_logger(__name__)
|
|
|
|
|
|
if is_tf_available():
|
|
import h5py
|
|
import numpy as np
|
|
import tensorflow as tf
|
|
|
|
from transformers import (
|
|
BertConfig,
|
|
PreTrainedModel,
|
|
PushToHubCallback,
|
|
RagRetriever,
|
|
TFAutoModel,
|
|
TFBertForMaskedLM,
|
|
TFBertForSequenceClassification,
|
|
TFBertModel,
|
|
TFPreTrainedModel,
|
|
TFRagModel,
|
|
)
|
|
from transformers.modeling_tf_utils import keras, tf_shard_checkpoint, unpack_inputs
|
|
from transformers.tf_utils import stable_softmax
|
|
|
|
tf.config.experimental.enable_tensor_float_32_execution(False)
|
|
|
|
if _tf_gpu_memory_limit is not None:
|
|
gpus = tf.config.list_physical_devices("GPU")
|
|
for gpu in gpus:
|
|
# Restrict TensorFlow to only allocate x GB of memory on the GPUs
|
|
try:
|
|
tf.config.set_logical_device_configuration(
|
|
gpu, [tf.config.LogicalDeviceConfiguration(memory_limit=_tf_gpu_memory_limit)]
|
|
)
|
|
logical_gpus = tf.config.list_logical_devices("GPU")
|
|
print("Logical GPUs", logical_gpus)
|
|
except RuntimeError as e:
|
|
# Virtual devices must be set before GPUs have been initialized
|
|
print(e)
|
|
|
|
if is_torch_available():
|
|
from transformers import BertModel
|
|
|
|
|
|
@require_tf
|
|
class TFModelUtilsTest(unittest.TestCase):
|
|
def test_cached_files_are_used_when_internet_is_down(self):
|
|
# A mock response for an HTTP head request to emulate server down
|
|
response_mock = mock.Mock()
|
|
response_mock.status_code = 500
|
|
response_mock.headers = {}
|
|
response_mock.raise_for_status.side_effect = HTTPError
|
|
response_mock.json.return_value = {}
|
|
|
|
# Download this model to make sure it's in the cache.
|
|
_ = TFBertModel.from_pretrained("hf-internal-testing/tiny-random-bert")
|
|
|
|
# Under the mock environment we get a 500 error when trying to reach the model.
|
|
with mock.patch("requests.Session.request", return_value=response_mock) as mock_head:
|
|
_ = TFBertModel.from_pretrained("hf-internal-testing/tiny-random-bert")
|
|
# This check we did call the fake head request
|
|
mock_head.assert_called()
|
|
|
|
# tests whether the unpack_inputs function behaves as expected
|
|
def test_unpack_inputs(self):
|
|
class DummyModel:
|
|
def __init__(self):
|
|
config_kwargs = {"output_attentions": False, "output_hidden_states": False, "return_dict": False}
|
|
self.config = PretrainedConfig(**config_kwargs)
|
|
self.main_input_name = "input_ids"
|
|
|
|
@unpack_inputs
|
|
def call(
|
|
self,
|
|
input_ids=None,
|
|
past_key_values=None,
|
|
output_attentions=None,
|
|
output_hidden_states=None,
|
|
return_dict=None,
|
|
):
|
|
return input_ids, past_key_values, output_attentions, output_hidden_states, return_dict
|
|
|
|
@unpack_inputs
|
|
def foo(self, pixel_values, output_attentions=None, output_hidden_states=None, return_dict=None):
|
|
return pixel_values, output_attentions, output_hidden_states, return_dict
|
|
|
|
dummy_model = DummyModel()
|
|
input_ids = tf.constant([0, 1, 2, 3], dtype=tf.int32)
|
|
past_key_values = tf.constant([4, 5, 6, 7], dtype=tf.int32)
|
|
pixel_values = tf.constant([8, 9, 10, 11], dtype=tf.int32)
|
|
|
|
# test case 1: Pass inputs as keyword arguments; Booleans are inherited from the config.
|
|
output = dummy_model.call(input_ids=input_ids, past_key_values=past_key_values)
|
|
tf.debugging.assert_equal(output[0], input_ids)
|
|
tf.debugging.assert_equal(output[1], past_key_values)
|
|
self.assertFalse(output[2])
|
|
self.assertFalse(output[3])
|
|
self.assertFalse(output[4])
|
|
|
|
# test case 2: Same as above, but with positional arguments.
|
|
output = dummy_model.call(input_ids, past_key_values)
|
|
tf.debugging.assert_equal(output[0], input_ids)
|
|
tf.debugging.assert_equal(output[1], past_key_values)
|
|
self.assertFalse(output[2])
|
|
self.assertFalse(output[3])
|
|
self.assertFalse(output[4])
|
|
|
|
# test case 3: We can also pack everything in the first input.
|
|
output = dummy_model.call(input_ids={"input_ids": input_ids, "past_key_values": past_key_values})
|
|
tf.debugging.assert_equal(output[0], input_ids)
|
|
tf.debugging.assert_equal(output[1], past_key_values)
|
|
self.assertFalse(output[2])
|
|
self.assertFalse(output[3])
|
|
self.assertFalse(output[4])
|
|
|
|
# test case 4: Explicit boolean arguments should override the config.
|
|
output = dummy_model.call(
|
|
input_ids=input_ids, past_key_values=past_key_values, output_attentions=False, return_dict=True
|
|
)
|
|
tf.debugging.assert_equal(output[0], input_ids)
|
|
tf.debugging.assert_equal(output[1], past_key_values)
|
|
self.assertFalse(output[2])
|
|
self.assertFalse(output[3])
|
|
self.assertTrue(output[4])
|
|
|
|
# test case 5: Unexpected arguments should raise an exception.
|
|
with self.assertRaises(ValueError):
|
|
output = dummy_model.call(input_ids=input_ids, past_key_values=past_key_values, foo="bar")
|
|
|
|
# test case 6: the decorator is independent from `main_input_name` -- it treats the first argument of the
|
|
# decorated function as its main input.
|
|
output = dummy_model.foo(pixel_values=pixel_values)
|
|
tf.debugging.assert_equal(output[0], pixel_values)
|
|
self.assertFalse(output[1])
|
|
self.assertFalse(output[2])
|
|
self.assertFalse(output[3])
|
|
|
|
# Tests whether the stable softmax is stable on CPU, with and without XLA
|
|
def test_xla_stable_softmax(self):
|
|
large_penalty = -1e9
|
|
n_tokens = 10
|
|
batch_size = 8
|
|
|
|
def masked_softmax(x, boolean_mask):
|
|
numerical_mask = (1.0 - tf.cast(boolean_mask, dtype=tf.float32)) * large_penalty
|
|
masked_x = x + numerical_mask
|
|
return stable_softmax(masked_x)
|
|
|
|
xla_masked_softmax = tf.function(masked_softmax, jit_compile=True)
|
|
xla_stable_softmax = tf.function(stable_softmax, jit_compile=True)
|
|
x = tf.random.normal((batch_size, n_tokens))
|
|
|
|
# Same outcome regardless of the boolean mask here
|
|
masked_tokens = random.randint(0, n_tokens)
|
|
boolean_mask = tf.convert_to_tensor([[1] * (n_tokens - masked_tokens) + [0] * masked_tokens], dtype=tf.int32)
|
|
|
|
# We can randomly mask a random numerical input OUTSIDE XLA
|
|
numerical_mask = (1.0 - tf.cast(boolean_mask, dtype=tf.float32)) * large_penalty
|
|
masked_x = x + numerical_mask
|
|
xla_out = xla_stable_softmax(masked_x)
|
|
out = stable_softmax(masked_x)
|
|
assert tf.experimental.numpy.allclose(xla_out, out)
|
|
|
|
# The stable softmax has the same output as the original softmax
|
|
unstable_out = tf.nn.softmax(masked_x)
|
|
assert tf.experimental.numpy.allclose(unstable_out, out)
|
|
|
|
# We can randomly mask a random numerical input INSIDE XLA
|
|
xla_out = xla_masked_softmax(x, boolean_mask)
|
|
out = masked_softmax(x, boolean_mask)
|
|
assert tf.experimental.numpy.allclose(xla_out, out)
|
|
|
|
def test_checkpoint_sharding_from_hub(self):
|
|
model = TFBertModel.from_pretrained("ArthurZ/tiny-random-bert-sharded")
|
|
# the model above is the same as the model below, just a sharded version.
|
|
ref_model = TFBertModel.from_pretrained("hf-internal-testing/tiny-random-bert")
|
|
for p1, p2 in zip(model.weights, ref_model.weights):
|
|
assert np.allclose(p1.numpy(), p2.numpy())
|
|
|
|
def test_sharded_checkpoint_with_prefix(self):
|
|
model = TFBertModel.from_pretrained("hf-internal-testing/tiny-random-bert", load_weight_prefix="a/b")
|
|
sharded_model = TFBertModel.from_pretrained("ArthurZ/tiny-random-bert-sharded", load_weight_prefix="a/b")
|
|
for p1, p2 in zip(model.weights, sharded_model.weights):
|
|
self.assertTrue(np.allclose(p1.numpy(), p2.numpy()))
|
|
self.assertTrue(p1.name.startswith("a/b/"))
|
|
self.assertTrue(p2.name.startswith("a/b/"))
|
|
|
|
def test_sharded_checkpoint_transfer(self):
|
|
# If this doesn't throw an error then the test passes
|
|
TFBertForSequenceClassification.from_pretrained("ArthurZ/tiny-random-bert-sharded")
|
|
|
|
@is_pt_tf_cross_test
|
|
def test_checkpoint_sharding_local_from_pt(self):
|
|
with tempfile.TemporaryDirectory() as tmp_dir:
|
|
_ = Repository(local_dir=tmp_dir, clone_from="hf-internal-testing/tiny-random-bert-sharded")
|
|
model = TFBertModel.from_pretrained(tmp_dir, from_pt=True)
|
|
# the model above is the same as the model below, just a sharded pytorch version.
|
|
ref_model = TFBertModel.from_pretrained("hf-internal-testing/tiny-random-bert")
|
|
for p1, p2 in zip(model.weights, ref_model.weights):
|
|
assert np.allclose(p1.numpy(), p2.numpy())
|
|
|
|
@is_pt_tf_cross_test
|
|
def test_checkpoint_loading_with_prefix_from_pt(self):
|
|
model = TFBertModel.from_pretrained(
|
|
"hf-internal-testing/tiny-random-bert", from_pt=True, load_weight_prefix="a/b"
|
|
)
|
|
ref_model = TFBertModel.from_pretrained("hf-internal-testing/tiny-random-bert", from_pt=True)
|
|
for p1, p2 in zip(model.weights, ref_model.weights):
|
|
self.assertTrue(np.allclose(p1.numpy(), p2.numpy()))
|
|
self.assertTrue(p1.name.startswith("a/b/"))
|
|
|
|
@is_pt_tf_cross_test
|
|
def test_checkpoint_sharding_hub_from_pt(self):
|
|
model = TFBertModel.from_pretrained("hf-internal-testing/tiny-random-bert-sharded", from_pt=True)
|
|
# the model above is the same as the model below, just a sharded pytorch version.
|
|
ref_model = TFBertModel.from_pretrained("hf-internal-testing/tiny-random-bert")
|
|
for p1, p2 in zip(model.weights, ref_model.weights):
|
|
assert np.allclose(p1.numpy(), p2.numpy())
|
|
|
|
def test_shard_checkpoint(self):
|
|
# This is the model we will use, total size 340,000 bytes.
|
|
model = keras.Sequential(
|
|
[
|
|
keras.layers.Dense(200, use_bias=False), # size 80,000
|
|
keras.layers.Dense(200, use_bias=False), # size 160,000
|
|
keras.layers.Dense(100, use_bias=False), # size 80,000
|
|
keras.layers.Dense(50, use_bias=False), # size 20,000
|
|
]
|
|
)
|
|
inputs = tf.zeros((1, 100), dtype=tf.float32)
|
|
model(inputs)
|
|
weights = model.weights
|
|
weights_dict = {w.name: w for w in weights}
|
|
with self.subTest("No shard when max size is bigger than model size"):
|
|
shards, index = tf_shard_checkpoint(weights)
|
|
self.assertIsNone(index)
|
|
self.assertDictEqual(shards, {TF2_WEIGHTS_NAME: weights})
|
|
|
|
with self.subTest("Test sharding, no weights bigger than max size"):
|
|
shards, index = tf_shard_checkpoint(weights, max_shard_size="300kB")
|
|
# Split is first two layers then last two.
|
|
self.assertDictEqual(
|
|
index,
|
|
{
|
|
"metadata": {"total_size": 340000},
|
|
"weight_map": {
|
|
"dense/kernel:0": "tf_model-00001-of-00002.h5",
|
|
"dense_1/kernel:0": "tf_model-00001-of-00002.h5",
|
|
"dense_2/kernel:0": "tf_model-00002-of-00002.h5",
|
|
"dense_3/kernel:0": "tf_model-00002-of-00002.h5",
|
|
},
|
|
},
|
|
)
|
|
|
|
shard1 = [weights_dict["dense/kernel:0"], weights_dict["dense_1/kernel:0"]]
|
|
shard2 = [weights_dict["dense_2/kernel:0"], weights_dict["dense_3/kernel:0"]]
|
|
self.assertDictEqual(shards, {"tf_model-00001-of-00002.h5": shard1, "tf_model-00002-of-00002.h5": shard2})
|
|
|
|
with self.subTest("Test sharding with weights bigger than max size"):
|
|
shards, index = tf_shard_checkpoint(weights, max_shard_size="100kB")
|
|
# Split is first layer, second layer then last 2.
|
|
self.assertDictEqual(
|
|
index,
|
|
{
|
|
"metadata": {"total_size": 340000},
|
|
"weight_map": {
|
|
"dense/kernel:0": "tf_model-00001-of-00003.h5",
|
|
"dense_1/kernel:0": "tf_model-00002-of-00003.h5",
|
|
"dense_2/kernel:0": "tf_model-00003-of-00003.h5",
|
|
"dense_3/kernel:0": "tf_model-00003-of-00003.h5",
|
|
},
|
|
},
|
|
)
|
|
|
|
shard1 = [weights_dict["dense/kernel:0"]]
|
|
shard2 = [weights_dict["dense_1/kernel:0"]]
|
|
shard3 = [weights_dict["dense_2/kernel:0"], weights_dict["dense_3/kernel:0"]]
|
|
self.assertDictEqual(
|
|
shards,
|
|
{
|
|
"tf_model-00001-of-00003.h5": shard1,
|
|
"tf_model-00002-of-00003.h5": shard2,
|
|
"tf_model-00003-of-00003.h5": shard3,
|
|
},
|
|
)
|
|
|
|
@slow
|
|
def test_special_layer_name_sharding(self):
|
|
retriever = RagRetriever.from_pretrained("facebook/rag-token-nq", index_name="exact", use_dummy_dataset=True)
|
|
model = TFRagModel.from_pretrained("facebook/rag-token-nq", retriever=retriever)
|
|
|
|
with tempfile.TemporaryDirectory() as tmp_dir:
|
|
for max_size in ["150kB", "150kiB", "200kB", "200kiB"]:
|
|
model.save_pretrained(tmp_dir, max_shard_size=max_size)
|
|
ref_model = TFRagModel.from_pretrained(tmp_dir, retriever=retriever)
|
|
for p1, p2 in zip(model.weights, ref_model.weights):
|
|
assert np.allclose(p1.numpy(), p2.numpy())
|
|
|
|
@require_safetensors
|
|
def test_checkpoint_sharding_local(self):
|
|
model = TFBertModel.from_pretrained("hf-internal-testing/tiny-random-bert")
|
|
|
|
with tempfile.TemporaryDirectory() as tmp_dir:
|
|
# We use the same folder for various sizes to make sure a new save erases the old checkpoint.
|
|
for max_size in ["150kB", "150kiB", "200kB", "200kiB"]:
|
|
model.save_pretrained(tmp_dir, max_shard_size=max_size)
|
|
|
|
# Get each shard file and its size
|
|
shard_to_size = {}
|
|
for shard in os.listdir(tmp_dir):
|
|
if shard.endswith(".h5"):
|
|
shard_file = os.path.join(tmp_dir, shard)
|
|
shard_to_size[shard_file] = os.path.getsize(shard_file)
|
|
|
|
index_file = os.path.join(tmp_dir, TF2_WEIGHTS_INDEX_NAME)
|
|
# Check there is an index but no regular weight file
|
|
self.assertTrue(os.path.isfile(index_file))
|
|
self.assertFalse(os.path.isfile(os.path.join(tmp_dir, TF2_WEIGHTS_NAME)))
|
|
|
|
# Check a file is bigger than max_size only when it has a single weight
|
|
for shard_file, size in shard_to_size.items():
|
|
if max_size.endswith("kiB"):
|
|
max_size_int = int(max_size[:-3]) * 2**10
|
|
else:
|
|
max_size_int = int(max_size[:-2]) * 10**3
|
|
# Note: pickle adds some junk so the weight of the file can end up being slightly bigger than
|
|
# the size asked for (since we count parameters)
|
|
if size >= max_size_int + 50000:
|
|
with h5py.File(shard_file, "r") as state_file:
|
|
self.assertEqual(len(state_file), 1)
|
|
|
|
# Check the index and the shard files found match
|
|
with open(index_file, "r", encoding="utf-8") as f:
|
|
index = json.loads(f.read())
|
|
|
|
all_shards = set(index["weight_map"].values())
|
|
shards_found = {f for f in os.listdir(tmp_dir) if f.endswith(".h5")}
|
|
self.assertSetEqual(all_shards, shards_found)
|
|
|
|
# Finally, check the model can be reloaded
|
|
new_model = TFBertModel.from_pretrained(tmp_dir)
|
|
|
|
model.build_in_name_scope()
|
|
new_model.build_in_name_scope()
|
|
|
|
for p1, p2 in zip(model.weights, new_model.weights):
|
|
self.assertTrue(np.allclose(p1.numpy(), p2.numpy()))
|
|
|
|
def test_safetensors_checkpoint_sharding_local(self):
|
|
model = TFBertModel.from_pretrained("hf-internal-testing/tiny-random-bert")
|
|
|
|
with tempfile.TemporaryDirectory() as tmp_dir:
|
|
# We use the same folder for various sizes to make sure a new save erases the old checkpoint.
|
|
for max_size in ["150kB", "150kiB", "200kB", "200kiB"]:
|
|
model.save_pretrained(tmp_dir, max_shard_size=max_size, safe_serialization=True)
|
|
|
|
# Get each shard file and its size
|
|
shard_to_size = {}
|
|
for shard in os.listdir(tmp_dir):
|
|
if shard.endswith(".h5"):
|
|
shard_file = os.path.join(tmp_dir, shard)
|
|
shard_to_size[shard_file] = os.path.getsize(shard_file)
|
|
|
|
index_file = os.path.join(tmp_dir, SAFE_WEIGHTS_INDEX_NAME)
|
|
# Check there is an index but no regular weight file
|
|
self.assertTrue(os.path.isfile(index_file))
|
|
self.assertFalse(os.path.isfile(os.path.join(tmp_dir, SAFE_WEIGHTS_NAME)))
|
|
self.assertFalse(os.path.isfile(os.path.join(tmp_dir, TF2_WEIGHTS_NAME)))
|
|
self.assertFalse(os.path.isfile(os.path.join(tmp_dir, TF2_WEIGHTS_INDEX_NAME)))
|
|
|
|
# Check the index and the shard files found match
|
|
with open(index_file, "r", encoding="utf-8") as f:
|
|
index = json.loads(f.read())
|
|
|
|
all_shards = set(index["weight_map"].values())
|
|
shards_found = {f for f in os.listdir(tmp_dir) if f.endswith(".safetensors")}
|
|
self.assertSetEqual(all_shards, shards_found)
|
|
|
|
# Finally, check the model can be reloaded
|
|
new_model = TFBertModel.from_pretrained(tmp_dir)
|
|
|
|
model.build_in_name_scope()
|
|
new_model.build_in_name_scope()
|
|
|
|
for p1, p2 in zip(model.weights, new_model.weights):
|
|
self.assertTrue(np.allclose(p1.numpy(), p2.numpy()))
|
|
|
|
@is_pt_tf_cross_test
|
|
@require_safetensors
|
|
def test_bfloat16_torch_loading(self):
|
|
# Assert that neither of these raise an error - both repos contain bfloat16 tensors
|
|
model1 = TFAutoModel.from_pretrained("Rocketknight1/tiny-random-gpt2-bfloat16-pt", from_pt=True)
|
|
model2 = TFAutoModel.from_pretrained("Rocketknight1/tiny-random-gpt2-bfloat16") # PT-format safetensors
|
|
# Check that PT and safetensors loading paths end up with the same values
|
|
for weight1, weight2 in zip(model1.weights, model2.weights):
|
|
self.assertTrue(tf.reduce_all(weight1 == weight2))
|
|
|
|
@slow
|
|
def test_save_pretrained_signatures(self):
|
|
model = TFBertModel.from_pretrained("hf-internal-testing/tiny-random-bert")
|
|
|
|
# Short custom TF signature function.
|
|
# `input_signature` is specific to BERT.
|
|
@tf.function(
|
|
input_signature=[
|
|
[
|
|
tf.TensorSpec([None, None], tf.int32, name="input_ids"),
|
|
tf.TensorSpec([None, None], tf.int32, name="token_type_ids"),
|
|
tf.TensorSpec([None, None], tf.int32, name="attention_mask"),
|
|
]
|
|
]
|
|
)
|
|
def serving_fn(input):
|
|
return model(input)
|
|
|
|
# Using default signature (default behavior) overrides 'serving_default'
|
|
with tempfile.TemporaryDirectory() as tmp_dir:
|
|
model.save_pretrained(tmp_dir, saved_model=True, signatures=None)
|
|
model_loaded = keras.models.load_model(f"{tmp_dir}/saved_model/1")
|
|
self.assertTrue("serving_default" in list(model_loaded.signatures.keys()))
|
|
|
|
# Providing custom signature function
|
|
with tempfile.TemporaryDirectory() as tmp_dir:
|
|
model.save_pretrained(tmp_dir, saved_model=True, signatures={"custom_signature": serving_fn})
|
|
model_loaded = keras.models.load_model(f"{tmp_dir}/saved_model/1")
|
|
self.assertTrue("custom_signature" in list(model_loaded.signatures.keys()))
|
|
|
|
# Providing multiple custom signature function
|
|
with tempfile.TemporaryDirectory() as tmp_dir:
|
|
model.save_pretrained(
|
|
tmp_dir,
|
|
saved_model=True,
|
|
signatures={"custom_signature_1": serving_fn, "custom_signature_2": serving_fn},
|
|
)
|
|
model_loaded = keras.models.load_model(f"{tmp_dir}/saved_model/1")
|
|
self.assertTrue("custom_signature_1" in list(model_loaded.signatures.keys()))
|
|
self.assertTrue("custom_signature_2" in list(model_loaded.signatures.keys()))
|
|
|
|
@require_safetensors
|
|
def test_safetensors_save_and_load(self):
|
|
model = TFBertModel.from_pretrained("hf-internal-testing/tiny-random-bert")
|
|
with tempfile.TemporaryDirectory() as tmp_dir:
|
|
model.save_pretrained(tmp_dir, safe_serialization=True)
|
|
# No tf_model.h5 file, only a model.safetensors
|
|
self.assertTrue(os.path.isfile(os.path.join(tmp_dir, SAFE_WEIGHTS_NAME)))
|
|
self.assertFalse(os.path.isfile(os.path.join(tmp_dir, SAFE_WEIGHTS_INDEX_NAME)))
|
|
self.assertFalse(os.path.isfile(os.path.join(tmp_dir, TF2_WEIGHTS_NAME)))
|
|
self.assertFalse(os.path.isfile(os.path.join(tmp_dir, TF2_WEIGHTS_INDEX_NAME)))
|
|
|
|
new_model = TFBertModel.from_pretrained(tmp_dir)
|
|
|
|
# Check models are equal
|
|
for p1, p2 in zip(model.weights, new_model.weights):
|
|
self.assertTrue(np.allclose(p1.numpy(), p2.numpy()))
|
|
|
|
@require_safetensors
|
|
def test_safetensors_sharded_save_and_load(self):
|
|
model = TFBertModel.from_pretrained("hf-internal-testing/tiny-random-bert")
|
|
with tempfile.TemporaryDirectory() as tmp_dir:
|
|
model.save_pretrained(tmp_dir, safe_serialization=True, max_shard_size="150kB")
|
|
# No tf weights or index file, only a safetensors index
|
|
self.assertFalse(os.path.isfile(os.path.join(tmp_dir, SAFE_WEIGHTS_NAME)))
|
|
self.assertFalse(os.path.isfile(os.path.join(tmp_dir, TF2_WEIGHTS_NAME)))
|
|
self.assertTrue(os.path.isfile(os.path.join(tmp_dir, SAFE_WEIGHTS_INDEX_NAME)))
|
|
self.assertFalse(os.path.isfile(os.path.join(tmp_dir, TF2_WEIGHTS_INDEX_NAME)))
|
|
|
|
new_model = TFBertModel.from_pretrained(tmp_dir)
|
|
|
|
# Check models are equal
|
|
for p1, p2 in zip(model.weights, new_model.weights):
|
|
self.assertTrue(np.allclose(p1.numpy(), p2.numpy()))
|
|
|
|
@is_pt_tf_cross_test
|
|
def test_safetensors_save_and_load_pt_to_tf(self):
|
|
model = TFBertModel.from_pretrained("hf-internal-testing/tiny-random-bert")
|
|
pt_model = BertModel.from_pretrained("hf-internal-testing/tiny-random-bert")
|
|
with tempfile.TemporaryDirectory() as tmp_dir:
|
|
pt_model.save_pretrained(tmp_dir, safe_serialization=True)
|
|
# Check we have a model.safetensors file
|
|
self.assertTrue(os.path.isfile(os.path.join(tmp_dir, SAFE_WEIGHTS_NAME)))
|
|
|
|
new_model = TFBertModel.from_pretrained(tmp_dir)
|
|
|
|
# Check models are equal
|
|
for p1, p2 in zip(model.weights, new_model.weights):
|
|
self.assertTrue(np.allclose(p1.numpy(), p2.numpy()))
|
|
|
|
@is_pt_tf_cross_test
|
|
def test_sharded_safetensors_save_and_load_pt_to_tf(self):
|
|
model = TFBertModel.from_pretrained("hf-internal-testing/tiny-random-bert")
|
|
pt_model = BertModel.from_pretrained("hf-internal-testing/tiny-random-bert")
|
|
with tempfile.TemporaryDirectory() as tmp_dir:
|
|
pt_model.save_pretrained(tmp_dir, safe_serialization=True, max_shard_size="150kB")
|
|
# Check we have a safetensors shard index file
|
|
self.assertTrue(os.path.isfile(os.path.join(tmp_dir, SAFE_WEIGHTS_INDEX_NAME)))
|
|
|
|
new_model = TFBertModel.from_pretrained(tmp_dir)
|
|
|
|
# Check models are equal
|
|
for p1, p2 in zip(model.weights, new_model.weights):
|
|
self.assertTrue(np.allclose(p1.numpy(), p2.numpy()))
|
|
|
|
@require_safetensors
|
|
def test_safetensors_load_from_hub(self):
|
|
tf_model = TFBertModel.from_pretrained("hf-internal-testing/tiny-random-bert")
|
|
|
|
# Can load from the TF-formatted checkpoint
|
|
safetensors_model = TFBertModel.from_pretrained("hf-internal-testing/tiny-random-bert-safetensors-tf")
|
|
|
|
# Check models are equal
|
|
for p1, p2 in zip(safetensors_model.weights, tf_model.weights):
|
|
self.assertTrue(np.allclose(p1.numpy(), p2.numpy()))
|
|
|
|
# Can load from the PyTorch-formatted checkpoint
|
|
safetensors_model = TFBertModel.from_pretrained("hf-internal-testing/tiny-random-bert-safetensors")
|
|
|
|
# Check models are equal
|
|
for p1, p2 in zip(safetensors_model.weights, tf_model.weights):
|
|
self.assertTrue(np.allclose(p1.numpy(), p2.numpy()))
|
|
|
|
@require_safetensors
|
|
def test_safetensors_tf_from_tf(self):
|
|
model = TFBertModel.from_pretrained("hf-internal-testing/tiny-bert-tf-only")
|
|
|
|
with tempfile.TemporaryDirectory() as tmp_dir:
|
|
model.save_pretrained(tmp_dir, safe_serialization=True)
|
|
new_model = TFBertModel.from_pretrained(tmp_dir)
|
|
|
|
for p1, p2 in zip(model.weights, new_model.weights):
|
|
self.assertTrue(np.allclose(p1.numpy(), p2.numpy()))
|
|
|
|
@require_safetensors
|
|
@is_pt_tf_cross_test
|
|
def test_safetensors_tf_from_torch(self):
|
|
hub_model = TFBertModel.from_pretrained("hf-internal-testing/tiny-bert-tf-only")
|
|
model = BertModel.from_pretrained("hf-internal-testing/tiny-bert-pt-only")
|
|
|
|
with tempfile.TemporaryDirectory() as tmp_dir:
|
|
model.save_pretrained(tmp_dir, safe_serialization=True)
|
|
new_model = TFBertModel.from_pretrained(tmp_dir)
|
|
|
|
for p1, p2 in zip(hub_model.weights, new_model.weights):
|
|
self.assertTrue(np.allclose(p1.numpy(), p2.numpy()))
|
|
|
|
@require_safetensors
|
|
def test_safetensors_tf_from_sharded_h5_with_sharded_safetensors_local(self):
|
|
with tempfile.TemporaryDirectory() as tmp_dir:
|
|
path = snapshot_download("hf-internal-testing/tiny-bert-tf-safetensors-h5-sharded", cache_dir=tmp_dir)
|
|
|
|
# This should not raise even if there are two types of sharded weights
|
|
TFBertModel.from_pretrained(path)
|
|
|
|
@require_safetensors
|
|
def test_safetensors_tf_from_sharded_h5_with_sharded_safetensors_hub(self):
|
|
# Confirm that we can correctly load the safetensors weights from a sharded hub repo even when TF weights present
|
|
TFBertModel.from_pretrained("hf-internal-testing/tiny-bert-tf-safetensors-h5-sharded", use_safetensors=True)
|
|
# Confirm that we can access the TF weights too
|
|
TFBertModel.from_pretrained("hf-internal-testing/tiny-bert-tf-safetensors-h5-sharded", use_safetensors=False)
|
|
|
|
@require_safetensors
|
|
def test_safetensors_load_from_local(self):
|
|
"""
|
|
This test checks that we can load safetensors from a checkpoint that only has those on the Hub
|
|
"""
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
location = snapshot_download("hf-internal-testing/tiny-bert-tf-only", cache_dir=tmp)
|
|
tf_model = TFBertModel.from_pretrained(location)
|
|
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
location = snapshot_download("hf-internal-testing/tiny-bert-tf-safetensors-only", cache_dir=tmp)
|
|
safetensors_model = TFBertModel.from_pretrained(location)
|
|
|
|
for p1, p2 in zip(tf_model.weights, safetensors_model.weights):
|
|
self.assertTrue(np.allclose(p1.numpy(), p2.numpy()))
|
|
|
|
@require_safetensors
|
|
def test_safetensors_load_from_hub_from_safetensors_pt(self):
|
|
"""
|
|
This test checks that we can load safetensors from a checkpoint that only has those on the Hub.
|
|
saved in the "pt" format.
|
|
"""
|
|
tf_model = TFBertModel.from_pretrained("hf-internal-testing/tiny-bert-h5")
|
|
|
|
# Can load from the PyTorch-formatted checkpoint
|
|
safetensors_model = TFBertModel.from_pretrained("hf-internal-testing/tiny-bert-pt-safetensors")
|
|
for p1, p2 in zip(tf_model.weights, safetensors_model.weights):
|
|
self.assertTrue(np.allclose(p1.numpy(), p2.numpy()))
|
|
|
|
@require_safetensors
|
|
def test_safetensors_load_from_local_from_safetensors_pt(self):
|
|
"""
|
|
This test checks that we can load safetensors from a local checkpoint that only has those
|
|
saved in the "pt" format.
|
|
"""
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
location = snapshot_download("hf-internal-testing/tiny-bert-h5", cache_dir=tmp)
|
|
tf_model = TFBertModel.from_pretrained(location)
|
|
|
|
# Can load from the PyTorch-formatted checkpoint
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
location = snapshot_download("hf-internal-testing/tiny-bert-pt-safetensors", cache_dir=tmp)
|
|
safetensors_model = TFBertModel.from_pretrained(location)
|
|
|
|
for p1, p2 in zip(tf_model.weights, safetensors_model.weights):
|
|
self.assertTrue(np.allclose(p1.numpy(), p2.numpy()))
|
|
|
|
@require_safetensors
|
|
def test_safetensors_load_from_hub_h5_before_safetensors(self):
|
|
"""
|
|
This test checks that we'll first download h5 weights before safetensors
|
|
The safetensors file on that repo is a pt safetensors and therefore cannot be loaded without PyTorch
|
|
"""
|
|
TFBertModel.from_pretrained("hf-internal-testing/tiny-bert-pt-safetensors-msgpack")
|
|
|
|
@require_safetensors
|
|
def test_safetensors_load_from_local_h5_before_safetensors(self):
|
|
"""
|
|
This test checks that we'll first download h5 weights before safetensors
|
|
The safetensors file on that repo is a pt safetensors and therefore cannot be loaded without PyTorch
|
|
"""
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
location = snapshot_download("hf-internal-testing/tiny-bert-pt-safetensors-msgpack", cache_dir=tmp)
|
|
TFBertModel.from_pretrained(location)
|
|
|
|
|
|
@require_tf
|
|
@is_staging_test
|
|
class TFModelPushToHubTester(unittest.TestCase):
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
cls._token = TOKEN
|
|
HfFolder.save_token(TOKEN)
|
|
|
|
@classmethod
|
|
def tearDownClass(cls):
|
|
try:
|
|
delete_repo(token=cls._token, repo_id="test-model-tf")
|
|
except HTTPError:
|
|
pass
|
|
|
|
try:
|
|
delete_repo(token=cls._token, repo_id="test-model-tf-callback")
|
|
except HTTPError:
|
|
pass
|
|
|
|
try:
|
|
delete_repo(token=cls._token, repo_id="valid_org/test-model-tf-org")
|
|
except HTTPError:
|
|
pass
|
|
|
|
def test_push_to_hub(self):
|
|
config = BertConfig(
|
|
vocab_size=99, hidden_size=32, num_hidden_layers=5, num_attention_heads=4, intermediate_size=37
|
|
)
|
|
model = TFBertModel(config)
|
|
# Make sure model is properly initialized
|
|
model.build_in_name_scope()
|
|
|
|
logging.set_verbosity_info()
|
|
logger = logging.get_logger("transformers.utils.hub")
|
|
with CaptureLogger(logger) as cl:
|
|
model.push_to_hub("test-model-tf", token=self._token)
|
|
logging.set_verbosity_warning()
|
|
# Check the model card was created and uploaded.
|
|
self.assertIn("Uploading the following files to __DUMMY_TRANSFORMERS_USER__/test-model-tf", cl.out)
|
|
|
|
new_model = TFBertModel.from_pretrained(f"{USER}/test-model-tf")
|
|
models_equal = True
|
|
for p1, p2 in zip(model.weights, new_model.weights):
|
|
if not tf.math.reduce_all(p1 == p2):
|
|
models_equal = False
|
|
break
|
|
self.assertTrue(models_equal)
|
|
|
|
# Reset repo
|
|
delete_repo(token=self._token, repo_id="test-model-tf")
|
|
|
|
# Push to hub via save_pretrained
|
|
with tempfile.TemporaryDirectory() as tmp_dir:
|
|
model.save_pretrained(tmp_dir, repo_id="test-model-tf", push_to_hub=True, token=self._token)
|
|
|
|
new_model = TFBertModel.from_pretrained(f"{USER}/test-model-tf")
|
|
models_equal = True
|
|
for p1, p2 in zip(model.weights, new_model.weights):
|
|
if not tf.math.reduce_all(p1 == p2):
|
|
models_equal = False
|
|
break
|
|
self.assertTrue(models_equal)
|
|
|
|
@is_pt_tf_cross_test
|
|
def test_push_to_hub_callback(self):
|
|
config = BertConfig(
|
|
vocab_size=99, hidden_size=32, num_hidden_layers=5, num_attention_heads=4, intermediate_size=37
|
|
)
|
|
model = TFBertForMaskedLM(config)
|
|
model.compile()
|
|
|
|
with tempfile.TemporaryDirectory() as tmp_dir:
|
|
push_to_hub_callback = PushToHubCallback(
|
|
output_dir=tmp_dir,
|
|
hub_model_id="test-model-tf-callback",
|
|
hub_token=self._token,
|
|
)
|
|
model.fit(model.dummy_inputs, model.dummy_inputs, epochs=1, callbacks=[push_to_hub_callback])
|
|
|
|
new_model = TFBertForMaskedLM.from_pretrained(f"{USER}/test-model-tf-callback")
|
|
models_equal = True
|
|
for p1, p2 in zip(model.weights, new_model.weights):
|
|
if not tf.math.reduce_all(p1 == p2):
|
|
models_equal = False
|
|
break
|
|
self.assertTrue(models_equal)
|
|
|
|
tf_push_to_hub_params = dict(inspect.signature(TFPreTrainedModel.push_to_hub).parameters)
|
|
tf_push_to_hub_params.pop("base_model_card_args")
|
|
pt_push_to_hub_params = dict(inspect.signature(PreTrainedModel.push_to_hub).parameters)
|
|
pt_push_to_hub_params.pop("deprecated_kwargs")
|
|
self.assertDictEaual(tf_push_to_hub_params, pt_push_to_hub_params)
|
|
|
|
def test_push_to_hub_in_organization(self):
|
|
config = BertConfig(
|
|
vocab_size=99, hidden_size=32, num_hidden_layers=5, num_attention_heads=4, intermediate_size=37
|
|
)
|
|
model = TFBertModel(config)
|
|
# Make sure model is properly initialized
|
|
model.build_in_name_scope()
|
|
|
|
model.push_to_hub("valid_org/test-model-tf-org", token=self._token)
|
|
|
|
new_model = TFBertModel.from_pretrained("valid_org/test-model-tf-org")
|
|
models_equal = True
|
|
for p1, p2 in zip(model.weights, new_model.weights):
|
|
if not tf.math.reduce_all(p1 == p2):
|
|
models_equal = False
|
|
break
|
|
self.assertTrue(models_equal)
|
|
|
|
# Reset repo
|
|
delete_repo(token=self._token, repo_id="valid_org/test-model-tf-org")
|
|
|
|
# Push to hub via save_pretrained
|
|
with tempfile.TemporaryDirectory() as tmp_dir:
|
|
model.save_pretrained(tmp_dir, push_to_hub=True, token=self._token, repo_id="valid_org/test-model-tf-org")
|
|
|
|
new_model = TFBertModel.from_pretrained("valid_org/test-model-tf-org")
|
|
models_equal = True
|
|
for p1, p2 in zip(model.weights, new_model.weights):
|
|
if not tf.math.reduce_all(p1 == p2):
|
|
models_equal = False
|
|
break
|
|
self.assertTrue(models_equal)
|