chendl/compositional_test
1
1# Copyright 2020 The HuggingFace Team. All rights reserved.2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7# http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14 15import copy16import inspect17import json18import random19import tempfile20import unittest21from typing import List, Tuple22 23import numpy as np24from huggingface_hub import HfFolder, delete_repo25from requests.exceptions import HTTPError26 27import transformers28from transformers import BertConfig, is_flax_available, is_torch_available29from transformers.models.auto import get_values30from transformers.testing_utils import (31 TOKEN,32 USER,33 CaptureLogger,34 is_pt_flax_cross_test,35 is_staging_test,36 require_flax,37 torch_device,38)39from transformers.utils import CONFIG_NAME, GENERATION_CONFIG_NAME, logging40from transformers.utils.generic import ModelOutput41 42 43if is_flax_available():44 import os45 46 import jax47 import jax.numpy as jnp48 from flax.core.frozen_dict import FrozenDict, freeze, unfreeze49 from flax.serialization import from_bytes50 from flax.traverse_util import flatten_dict, unflatten_dict51 52 from transformers import (53 FLAX_MODEL_FOR_QUESTION_ANSWERING_MAPPING,54 FLAX_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING,55 FLAX_MODEL_MAPPING,56 FlaxAutoModel,57 FlaxAutoModelForSequenceClassification,58 FlaxBertModel,59 )60 from transformers.modeling_flax_pytorch_utils import (61 convert_pytorch_state_dict_to_flax,62 load_flax_weights_in_pytorch_model,63 )64 from transformers.modeling_flax_utils import FLAX_WEIGHTS_INDEX_NAME, FLAX_WEIGHTS_NAME65 66 os.environ["XLA_PYTHON_CLIENT_MEM_FRACTION"] = "0.12" # assumed parallelism: 867 68if is_torch_available():69 import torch70 71 72def _config_zero_init(config):73 configs_no_init = copy.deepcopy(config)74 for key in configs_no_init.__dict__.keys():75 if "_range" in key or "_std" in key or "initializer_factor" in key:76 setattr(configs_no_init, key, 1e-10)77 return configs_no_init78 79 80def ids_tensor(shape, vocab_size, rng=None):81 """Creates a random int32 tensor of the shape within the vocab size."""82 if rng is None:83 rng = random.Random()84 85 total_dims = 186 for dim in shape:87 total_dims *= dim88 89 values = []90 for _ in range(total_dims):91 values.append(rng.randint(0, vocab_size - 1))92 93 output = np.array(values, dtype=jnp.int32).reshape(shape)94 95 return output96 97 98def floats_tensor(shape, scale=1.0, rng=None, name=None):99 """Creates a random float32 tensor"""100 if rng is None:101 rng = random.Random()102 103 total_dims = 1104 for dim in shape:105 total_dims *= dim106 107 values = []108 for _ in range(total_dims):109 values.append(rng.random() * scale)110 111 return np.array(values, dtype=jnp.float32).reshape(shape)112 113 114def random_attention_mask(shape, rng=None):115 attn_mask = ids_tensor(shape, vocab_size=2, rng=rng)116 # make sure that at least one token is attended to for each batch117 attn_mask[:, -1] = 1118 return attn_mask119 120 121def get_params(params, from_head_prefix=None):122 """Function extracts relevant parameters into flatten dict from model params,123 appends batch normalization statistics if present"""124 125 # If Both parameters and batch normalization statistics are present126 if "batch_stats" in params:127 # Extract only parameters for the specified head prefix (if specified) and add batch statistics128 if from_head_prefix is not None:129 extracted_params = flatten_dict(unfreeze(params["params"][from_head_prefix]))130 extracted_params.update(flatten_dict(params["batch_stats"][from_head_prefix]))131 else:132 extracted_params = flatten_dict(unfreeze(params["params"]))133 extracted_params.update(flatten_dict(params["batch_stats"]))134 135 # Only parameters are present136 else:137 if from_head_prefix is not None:138 extracted_params = flatten_dict(unfreeze(params[from_head_prefix]))139 else:140 extracted_params = flatten_dict(unfreeze(params))141 142 return extracted_params143 144 145@require_flax146class FlaxModelTesterMixin:147 model_tester = None148 all_model_classes = ()149 test_mismatched_shapes = True150 is_encoder_decoder = False151 test_head_masking = False152 has_attentions = True153 154 def _prepare_for_class(self, inputs_dict, model_class):155 inputs_dict = copy.deepcopy(inputs_dict)156 157 # hack for now until we have AutoModel classes158 if "ForMultipleChoice" in model_class.__name__:159 inputs_dict = {160 k: jnp.broadcast_to(v[:, None], (v.shape[0], self.model_tester.num_choices, v.shape[-1]))161 if isinstance(v, (jnp.ndarray, np.ndarray))162 else v163 for k, v in inputs_dict.items()164 }165 166 return inputs_dict167 168 def assert_almost_equals(self, a: np.ndarray, b: np.ndarray, tol: float):169 diff = np.abs((a - b)).max()170 self.assertLessEqual(diff, tol, f"Difference between torch and flax is {diff} (>= {tol}).")171 172 def test_model_outputs_equivalence(self):173 config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()174 175 def check_equivalence(model, tuple_inputs, dict_inputs, additional_kwargs={}):176 tuple_output = model(**tuple_inputs, return_dict=False, **additional_kwargs)177 dict_output = model(**dict_inputs, return_dict=True, **additional_kwargs).to_tuple()178 179 def recursive_check(tuple_object, dict_object):180 if isinstance(tuple_object, (List, Tuple)):181 for tuple_iterable_value, dict_iterable_value in zip(tuple_object, dict_object):182 recursive_check(tuple_iterable_value, dict_iterable_value)183 elif tuple_object is None:184 return185 else:186 self.assert_almost_equals(jnp.nan_to_num(tuple_object), jnp.nan_to_num(dict_object), 1e-5)187 188 recursive_check(tuple_output, dict_output)189 190 for model_class in self.all_model_classes:191 model = model_class(config)192 193 tuple_inputs = self._prepare_for_class(inputs_dict, model_class)194 dict_inputs = self._prepare_for_class(inputs_dict, model_class)195 check_equivalence(model, tuple_inputs, dict_inputs)196 197 tuple_inputs = self._prepare_for_class(inputs_dict, model_class)198 dict_inputs = self._prepare_for_class(inputs_dict, model_class)199 check_equivalence(model, tuple_inputs, dict_inputs, {"output_hidden_states": True})200 201 # (Copied from tests.test_modeling_common.ModelTesterMixin.check_pt_flax_outputs)202 def check_pt_flax_outputs(self, fx_outputs, pt_outputs, model_class, tol=1e-5, name="outputs", attributes=None):203 """204 Args:205 model_class: The class of the model that is currently testing. For example, ..., etc.206 Currently unused, but it could make debugging easier and faster.207 208 names: A string, or a list of strings. These specify what fx_outputs/pt_outputs represent in the model outputs.209 Currently unused, but in the future, we could use this information to make the error message clearer210 by giving the name(s) of the output tensor(s) with large difference(s) between PT and Flax.211 """212 213 self.assertEqual(type(name), str)214 if attributes is not None:215 self.assertEqual(type(attributes), tuple, f"{name}: The argument `attributes` should be a `tuple`")216 217 # Allow `ModelOutput` (e.g. `CLIPOutput` has `text_model_output` and `vision_model_output`).218 if isinstance(fx_outputs, ModelOutput):219 self.assertTrue(220 isinstance(pt_outputs, ModelOutput),221 f"{name}: `pt_outputs` should an instance of `ModelOutput` when `fx_outputs` is",222 )223 224 fx_keys = tuple([k for k, v in fx_outputs.items() if v is not None])225 pt_keys = tuple([k for k, v in pt_outputs.items() if v is not None])226 227 self.assertEqual(fx_keys, pt_keys, f"{name}: Output keys differ between Flax and PyTorch")228 229 # convert to the case of `tuple`230 # appending each key to the current (string) `name`231 attributes = tuple([f"{name}.{k}" for k in fx_keys])232 self.check_pt_flax_outputs(233 fx_outputs.to_tuple(), pt_outputs.to_tuple(), model_class, tol=tol, name=name, attributes=attributes234 )235 236 # Allow `list` (e.g. `TransfoXLModelOutput.mems` is a list of tensors.)237 elif type(fx_outputs) in [tuple, list]:238 self.assertEqual(239 type(fx_outputs), type(pt_outputs), f"{name}: Output types differ between Flax and PyTorch"240 )241 self.assertEqual(242 len(fx_outputs), len(pt_outputs), f"{name}: Output lengths differ between Flax and PyTorch"243 )244 245 if attributes is not None:246 # case 1: each output has assigned name (e.g. a tuple form of a `ModelOutput`)247 self.assertEqual(248 len(attributes),249 len(fx_outputs),250 f"{name}: The tuple `attributes` should have the same length as `fx_outputs`",251 )252 else:253 # case 2: each output has no assigned name (e.g. hidden states of each layer) -> add an index to `name`254 attributes = tuple([f"{name}_{idx}" for idx in range(len(fx_outputs))])255 256 for fx_output, pt_output, attr in zip(fx_outputs, pt_outputs, attributes):257 self.check_pt_flax_outputs(fx_output, pt_output, model_class, tol=tol, name=attr)258 259 elif isinstance(fx_outputs, jnp.ndarray):260 self.assertTrue(261 isinstance(pt_outputs, torch.Tensor), f"{name}: `pt_outputs` should a tensor when `fx_outputs` is"262 )263 264 # Using `np.asarray` gives `ValueError: assignment destination is read-only` at the line `fx_outputs[fx_nans] = 0`.265 fx_outputs = np.array(fx_outputs)266 pt_outputs = pt_outputs.detach().to("cpu").numpy()267 268 self.assertEqual(269 fx_outputs.shape, pt_outputs.shape, f"{name}: Output shapes differ between Flax and PyTorch"270 )271 272 # deal with NumPy's scalars to make replacing nan values by 0 work.273 if np.isscalar(fx_outputs):274 fx_outputs = np.array([fx_outputs])275 pt_outputs = np.array([pt_outputs])276 277 fx_nans = np.isnan(fx_outputs)278 pt_nans = np.isnan(pt_outputs)279 280 pt_outputs[fx_nans] = 0281 fx_outputs[fx_nans] = 0282 pt_outputs[pt_nans] = 0283 fx_outputs[pt_nans] = 0284 285 max_diff = np.amax(np.abs(fx_outputs - pt_outputs))286 self.assertLessEqual(287 max_diff, tol, f"{name}: Difference between PyTorch and Flax is {max_diff} (>= {tol})."288 )289 else:290 raise ValueError(291 "`fx_outputs` should be an instance of `ModelOutput`, a `tuple`, or an instance of `jnp.ndarray`. Got"292 f" {type(fx_outputs)} instead."293 )294 295 @is_pt_flax_cross_test296 def test_equivalence_pt_to_flax(self):297 # It might be better to put this inside the for loop below (because we modify the config there).298 # But logically, it is fine.299 config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()300 301 for model_class in self.all_model_classes:302 with self.subTest(model_class.__name__):303 # Output all for aggressive testing304 config.output_hidden_states = True305 config.output_attentions = self.has_attentions306 307 # prepare inputs308 prepared_inputs_dict = self._prepare_for_class(inputs_dict, model_class)309 pt_inputs = {k: torch.tensor(v.tolist(), device=torch_device) for k, v in prepared_inputs_dict.items()}310 311 # load corresponding PyTorch class312 pt_model_class_name = model_class.__name__[4:] # Skip the "Flax" at the beginning313 pt_model_class = getattr(transformers, pt_model_class_name)314 315 pt_model = pt_model_class(config).eval()316 # Flax models don't use the `use_cache` option and cache is not returned as a default.317 # So we disable `use_cache` here for PyTorch model.318 pt_model.config.use_cache = False319 fx_model = model_class(config, dtype=jnp.float32)320 321 fx_state = convert_pytorch_state_dict_to_flax(pt_model.state_dict(), fx_model)322 fx_model.params = fx_state323 324 # send pytorch model to the correct device325 pt_model.to(torch_device)326 327 with torch.no_grad():328 pt_outputs = pt_model(**pt_inputs)329 fx_outputs = fx_model(**prepared_inputs_dict)330 331 fx_keys = tuple([k for k, v in fx_outputs.items() if v is not None])332 pt_keys = tuple([k for k, v in pt_outputs.items() if v is not None])333 334 self.assertEqual(fx_keys, pt_keys)335 self.check_pt_flax_outputs(fx_outputs, pt_outputs, model_class)336 337 with tempfile.TemporaryDirectory() as tmpdirname:338 pt_model.save_pretrained(tmpdirname)339 fx_model_loaded = model_class.from_pretrained(tmpdirname, from_pt=True)340 341 fx_outputs_loaded = fx_model_loaded(**prepared_inputs_dict)342 343 fx_keys = tuple([k for k, v in fx_outputs_loaded.items() if v is not None])344 pt_keys = tuple([k for k, v in pt_outputs.items() if v is not None])345 346 self.assertEqual(fx_keys, pt_keys)347 self.check_pt_flax_outputs(fx_outputs_loaded, pt_outputs, model_class)348 349 @is_pt_flax_cross_test350 def test_equivalence_flax_to_pt(self):351 config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()352 353 for model_class in self.all_model_classes:354 with self.subTest(model_class.__name__):355 # Output all for aggressive testing356 config.output_hidden_states = True357 config.output_attentions = self.has_attentions358 359 # prepare inputs360 prepared_inputs_dict = self._prepare_for_class(inputs_dict, model_class)361 pt_inputs = {k: torch.tensor(v.tolist(), device=torch_device) for k, v in prepared_inputs_dict.items()}362 363 # load corresponding PyTorch class364 pt_model_class_name = model_class.__name__[4:] # Skip the "Flax" at the beginning365 pt_model_class = getattr(transformers, pt_model_class_name)366 367 pt_model = pt_model_class(config).eval()368 # Flax models don't use the `use_cache` option and cache is not returned as a default.369 # So we disable `use_cache` here for PyTorch model.370 pt_model.config.use_cache = False371 fx_model = model_class(config, dtype=jnp.float32)372 373 pt_model = load_flax_weights_in_pytorch_model(pt_model, fx_model.params)374 375 # make sure weights are tied in PyTorch376 pt_model.tie_weights()377 378 # send pytorch model to the correct device379 pt_model.to(torch_device)380 381 with torch.no_grad():382 pt_outputs = pt_model(**pt_inputs)383 fx_outputs = fx_model(**prepared_inputs_dict)384 385 fx_keys = tuple([k for k, v in fx_outputs.items() if v is not None])386 pt_keys = tuple([k for k, v in pt_outputs.items() if v is not None])387 388 self.assertEqual(fx_keys, pt_keys)389 self.check_pt_flax_outputs(fx_outputs, pt_outputs, model_class)390 391 with tempfile.TemporaryDirectory() as tmpdirname:392 fx_model.save_pretrained(tmpdirname)393 pt_model_loaded = pt_model_class.from_pretrained(tmpdirname, from_flax=True)394 395 # send pytorch model to the correct device396 pt_model_loaded.to(torch_device)397 pt_model_loaded.eval()398 399 with torch.no_grad():400 pt_outputs_loaded = pt_model_loaded(**pt_inputs)401 402 fx_keys = tuple([k for k, v in fx_outputs.items() if v is not None])403 pt_keys = tuple([k for k, v in pt_outputs_loaded.items() if v is not None])404 405 self.assertEqual(fx_keys, pt_keys)406 self.check_pt_flax_outputs(fx_outputs, pt_outputs_loaded, model_class)407 408 def test_from_pretrained_save_pretrained(self):409 config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()410 411 for model_class in self.all_model_classes:412 with self.subTest(model_class.__name__):413 model = model_class(config)414 415 prepared_inputs_dict = self._prepare_for_class(inputs_dict, model_class)416 outputs = model(**prepared_inputs_dict).to_tuple()417 418 # verify that normal save_pretrained works as expected419 with tempfile.TemporaryDirectory() as tmpdirname:420 model.save_pretrained(tmpdirname)421 422 # the config file (and the generation config file, if it can generate) should be saved423 self.assertTrue(os.path.exists(os.path.join(tmpdirname, CONFIG_NAME)))424 self.assertEqual(425 model.can_generate(), os.path.exists(os.path.join(tmpdirname, GENERATION_CONFIG_NAME))426 )427 428 model_loaded = model_class.from_pretrained(tmpdirname)429 430 outputs_loaded = model_loaded(**prepared_inputs_dict).to_tuple()431 for output_loaded, output in zip(outputs_loaded, outputs):432 self.assert_almost_equals(output_loaded, output, 1e-3)433 434 # verify that save_pretrained for distributed training435 # with `params=params` works as expected436 with tempfile.TemporaryDirectory() as tmpdirname:437 model.save_pretrained(tmpdirname, params=model.params)438 model_loaded = model_class.from_pretrained(tmpdirname)439 440 outputs_loaded = model_loaded(**prepared_inputs_dict).to_tuple()441 for output_loaded, output in zip(outputs_loaded, outputs):442 self.assert_almost_equals(output_loaded, output, 1e-3)443 444 def test_save_load_from_base(self):445 config, _ = self.model_tester.prepare_config_and_inputs_for_common()446 base_class = FLAX_MODEL_MAPPING[config.__class__]447 448 for model_class in self.all_model_classes:449 if model_class == base_class:450 continue451 452 model = base_class(config)453 base_params = get_params(model.params)454 455 # check that all base model weights are loaded correctly456 with tempfile.TemporaryDirectory() as tmpdirname:457 model.save_pretrained(tmpdirname)458 head_model = model_class.from_pretrained(tmpdirname)459 460 base_param_from_head = get_params(head_model.params, from_head_prefix=head_model.base_model_prefix)461 462 for key in base_param_from_head.keys():463 max_diff = (base_params[key] - base_param_from_head[key]).sum().item()464 self.assertLessEqual(max_diff, 1e-3, msg=f"{key} not identical")465 466 def test_save_load_to_base(self):467 config, _ = self.model_tester.prepare_config_and_inputs_for_common()468 base_class = FLAX_MODEL_MAPPING[config.__class__]469 470 for model_class in self.all_model_classes:471 if model_class == base_class:472 continue473 474 model = model_class(config)475 base_params_from_head = get_params(model.params, from_head_prefix=model.base_model_prefix)476 477 # check that all base model weights are loaded correctly478 with tempfile.TemporaryDirectory() as tmpdirname:479 model.save_pretrained(tmpdirname)480 base_model = base_class.from_pretrained(tmpdirname)481 482 base_params = get_params(base_model.params)483 484 for key in base_params_from_head.keys():485 max_diff = (base_params[key] - base_params_from_head[key]).sum().item()486 self.assertLessEqual(max_diff, 1e-3, msg=f"{key} not identical")487 488 @is_pt_flax_cross_test489 def test_save_load_from_base_pt(self):490 config, _ = self.model_tester.prepare_config_and_inputs_for_common()491 base_class = FLAX_MODEL_MAPPING[config.__class__]492 493 for model_class in self.all_model_classes:494 if model_class == base_class:495 continue496 497 model = base_class(config)498 base_params = get_params(model.params)499 500 # convert Flax model to PyTorch model501 pt_model_class = getattr(transformers, base_class.__name__[4:]) # Skip the "Flax" at the beginning502 pt_model = pt_model_class(config).eval()503 pt_model = load_flax_weights_in_pytorch_model(pt_model, model.params)504 505 # check that all base model weights are loaded correctly506 with tempfile.TemporaryDirectory() as tmpdirname:507 # save pt model508 pt_model.save_pretrained(tmpdirname)509 head_model = model_class.from_pretrained(tmpdirname, from_pt=True)510 511 base_param_from_head = get_params(head_model.params, from_head_prefix=head_model.base_model_prefix)512 513 for key in base_param_from_head.keys():514 max_diff = (base_params[key] - base_param_from_head[key]).sum().item()515 self.assertLessEqual(max_diff, 1e-3, msg=f"{key} not identical")516 517 @is_pt_flax_cross_test518 def test_save_load_to_base_pt(self):519 config, _ = self.model_tester.prepare_config_and_inputs_for_common()520 base_class = FLAX_MODEL_MAPPING[config.__class__]521 522 for model_class in self.all_model_classes:523 if model_class == base_class:524 continue525 526 model = model_class(config)527 base_params_from_head = get_params(model.params, from_head_prefix=model.base_model_prefix)528 529 # convert Flax model to PyTorch model530 pt_model_class = getattr(transformers, model_class.__name__[4:]) # Skip the "Flax" at the beginning531 pt_model = pt_model_class(config).eval()532 pt_model = load_flax_weights_in_pytorch_model(pt_model, model.params)533 534 # check that all base model weights are loaded correctly535 with tempfile.TemporaryDirectory() as tmpdirname:536 pt_model.save_pretrained(tmpdirname)537 base_model = base_class.from_pretrained(tmpdirname, from_pt=True)538 539 base_params = get_params(base_model.params)540 541 for key in base_params_from_head.keys():542 max_diff = (base_params[key] - base_params_from_head[key]).sum().item()543 self.assertLessEqual(max_diff, 1e-3, msg=f"{key} not identical")544 545 @is_pt_flax_cross_test546 def test_save_load_bf16_to_base_pt(self):547 config, _ = self.model_tester.prepare_config_and_inputs_for_common()548 base_class = FLAX_MODEL_MAPPING[config.__class__]549 550 for model_class in self.all_model_classes:551 if model_class == base_class:552 continue553 554 model = model_class(config)555 model.params = model.to_bf16(model.params)556 base_params_from_head = get_params(model.params, from_head_prefix=model.base_model_prefix)557 558 # convert Flax model to PyTorch model559 pt_model_class = getattr(transformers, model_class.__name__[4:]) # Skip the "Flax" at the beginning560 pt_model = pt_model_class(config).eval()561 pt_model = load_flax_weights_in_pytorch_model(pt_model, model.params)562 563 # check that all base model weights are loaded correctly564 with tempfile.TemporaryDirectory() as tmpdirname:565 pt_model.save_pretrained(tmpdirname)566 base_model = base_class.from_pretrained(tmpdirname, from_pt=True)567 568 base_params = get_params(base_model.params)569 570 for key in base_params_from_head.keys():571 max_diff = (base_params[key] - base_params_from_head[key]).sum().item()572 self.assertLessEqual(max_diff, 1e-3, msg=f"{key} not identical")573 574 def test_jit_compilation(self):575 config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()576 577 for model_class in self.all_model_classes:578 with self.subTest(model_class.__name__):579 prepared_inputs_dict = self._prepare_for_class(inputs_dict, model_class)580 model = model_class(config)581 582 @jax.jit583 def model_jitted(input_ids, attention_mask=None, **kwargs):584 return model(input_ids=input_ids, attention_mask=attention_mask, **kwargs)585 586 with self.subTest("JIT Enabled"):587 jitted_outputs = model_jitted(**prepared_inputs_dict).to_tuple()588 589 with self.subTest("JIT Disabled"):590 with jax.disable_jit():591 outputs = model_jitted(**prepared_inputs_dict).to_tuple()592 593 self.assertEqual(len(outputs), len(jitted_outputs))594 for jitted_output, output in zip(jitted_outputs, outputs):595 self.assertEqual(jitted_output.shape, output.shape)596 597 def test_forward_signature(self):598 config, _ = self.model_tester.prepare_config_and_inputs_for_common()599 600 for model_class in self.all_model_classes:601 model = model_class(config)602 signature = inspect.signature(model.__call__)603 # signature.parameters is an OrderedDict => so arg_names order is deterministic604 arg_names = [*signature.parameters.keys()]605 606 if model.config.is_encoder_decoder:607 expected_arg_names = [608 "input_ids",609 "attention_mask",610 "decoder_input_ids",611 "decoder_attention_mask",612 ]613 self.assertListEqual(arg_names[: len(expected_arg_names)], expected_arg_names)614 else:615 expected_arg_names = ["input_ids", "attention_mask"]616 self.assertListEqual(arg_names[:2], expected_arg_names)617 618 def test_naming_convention(self):619 for model_class in self.all_model_classes:620 model_class_name = model_class.__name__621 module_class_name = (622 model_class_name[:-5] + "Module" if model_class_name[-5:] == "Model" else model_class_name + "Module"623 )624 bert_modeling_flax_module = __import__(model_class.__module__, fromlist=[module_class_name])625 module_cls = getattr(bert_modeling_flax_module, module_class_name)626 627 self.assertIsNotNone(module_cls)628 629 def test_hidden_states_output(self):630 def check_hidden_states_output(inputs_dict, config, model_class):631 model = model_class(config)632 633 outputs = model(**self._prepare_for_class(inputs_dict, model_class))634 hidden_states = outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states635 636 expected_num_layers = getattr(637 self.model_tester, "expected_num_hidden_layers", self.model_tester.num_hidden_layers + 1638 )639 self.assertEqual(len(hidden_states), expected_num_layers)640 641 if hasattr(self.model_tester, "encoder_seq_length"):642 seq_length = self.model_tester.encoder_seq_length643 else:644 seq_length = self.model_tester.seq_length645 646 self.assertListEqual(647 list(hidden_states[0].shape[-2:]),648 [seq_length, self.model_tester.hidden_size],649 )650 651 if config.is_encoder_decoder:652 hidden_states = outputs.decoder_hidden_states653 654 self.assertIsInstance(hidden_states, (list, tuple))655 self.assertEqual(len(hidden_states), expected_num_layers)656 seq_len = getattr(self.model_tester, "seq_length", None)657 decoder_seq_length = getattr(self.model_tester, "decoder_seq_length", seq_len)658 659 self.assertListEqual(660 list(hidden_states[0].shape[-2:]),661 [decoder_seq_length, self.model_tester.hidden_size],662 )663 664 config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()665 666 for model_class in self.all_model_classes:667 inputs_dict["output_hidden_states"] = True668 check_hidden_states_output(inputs_dict, config, model_class)669 670 # check that output_hidden_states also work using config671 del inputs_dict["output_hidden_states"]672 config.output_hidden_states = True673 674 check_hidden_states_output(inputs_dict, config, model_class)675 676 def test_attention_outputs(self):677 if not self.has_attentions:678 self.skipTest(reason="Model does not output attentions")679 680 config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()681 config.return_dict = True682 683 seq_length = getattr(self.model_tester, "seq_length", None)684 decoder_seq_length = getattr(self.model_tester, "decoder_seq_length", seq_length)685 encoder_seq_length = getattr(self.model_tester, "encoder_seq_length", seq_length)686 decoder_key_length = getattr(self.model_tester, "decoder_key_length", decoder_seq_length)687 encoder_key_length = getattr(self.model_tester, "key_length", encoder_seq_length)688 689 for model_class in self.all_model_classes:690 inputs_dict["output_attentions"] = True691 inputs_dict["output_hidden_states"] = False692 model = model_class(config)693 outputs = model(**self._prepare_for_class(inputs_dict, model_class))694 attentions = outputs.encoder_attentions if config.is_encoder_decoder else outputs.attentions695 self.assertEqual(len(attentions), self.model_tester.num_hidden_layers)696 697 # check that output_attentions also work using config698 del inputs_dict["output_attentions"]699 config.output_attentions = True700 model = model_class(config)701 outputs = model(**self._prepare_for_class(inputs_dict, model_class))702 attentions = outputs.encoder_attentions if config.is_encoder_decoder else outputs.attentions703 self.assertEqual(len(attentions), self.model_tester.num_hidden_layers)704 705 self.assertListEqual(706 list(attentions[0].shape[-3:]),707 [self.model_tester.num_attention_heads, encoder_seq_length, encoder_key_length],708 )709 out_len = len(outputs)710 711 if self.is_encoder_decoder:712 correct_outlen = 5713 714 # Question Answering model returns start_logits and end_logits715 if model_class in get_values(FLAX_MODEL_FOR_QUESTION_ANSWERING_MAPPING):716 correct_outlen += 1 # start_logits and end_logits instead of only 1 output717 718 self.assertEqual(out_len, correct_outlen)719 720 # decoder attentions721 decoder_attentions = outputs.decoder_attentions722 self.assertIsInstance(decoder_attentions, (list, tuple))723 self.assertEqual(len(decoder_attentions), self.model_tester.num_hidden_layers)724 self.assertListEqual(725 list(decoder_attentions[0].shape[-3:]),726 [self.model_tester.num_attention_heads, decoder_seq_length, decoder_key_length],727 )728 729 # cross attentions730 cross_attentions = outputs.cross_attentions731 self.assertIsInstance(cross_attentions, (list, tuple))732 self.assertEqual(len(cross_attentions), self.model_tester.num_hidden_layers)733 self.assertListEqual(734 list(cross_attentions[0].shape[-3:]),735 [736 self.model_tester.num_attention_heads,737 decoder_seq_length,738 encoder_key_length,739 ],740 )741 742 # Check attention is always last and order is fine743 inputs_dict["output_attentions"] = True744 inputs_dict["output_hidden_states"] = True745 model = model_class(config)746 outputs = model(**self._prepare_for_class(inputs_dict, model_class))747 748 if hasattr(self.model_tester, "num_hidden_states_types"):749 added_hidden_states = self.model_tester.num_hidden_states_types750 elif self.is_encoder_decoder:751 added_hidden_states = 2752 else:753 added_hidden_states = 1754 self.assertEqual(out_len + added_hidden_states, len(outputs))755 756 self_attentions = outputs.encoder_attentions if config.is_encoder_decoder else outputs.attentions757 self.assertEqual(len(self_attentions), self.model_tester.num_hidden_layers)758 759 self.assertListEqual(760 list(self_attentions[0].shape[-3:]),761 [self.model_tester.num_attention_heads, encoder_seq_length, encoder_key_length],762 )763 764 def test_load_with_mismatched_shapes(self):765 if not self.test_mismatched_shapes:766 return767 config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()768 769 for model_class in self.all_model_classes:770 if model_class not in get_values(FLAX_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING):771 continue772 773 with self.subTest(msg=f"Testing {model_class}"):774 with tempfile.TemporaryDirectory() as tmp_dir:775 model = model_class(config)776 model.save_pretrained(tmp_dir)777 778 # Fails when we don't set ignore_mismatched_sizes=True779 with self.assertRaises(ValueError):780 new_model = FlaxAutoModelForSequenceClassification.from_pretrained(tmp_dir, num_labels=42)781 with self.assertRaises(ValueError):782 new_model_without_prefix = FlaxAutoModel.from_pretrained(tmp_dir, vocab_size=10)783 784 logger = logging.get_logger("transformers.modeling_flax_utils")785 with CaptureLogger(logger) as cl:786 new_model = FlaxAutoModelForSequenceClassification.from_pretrained(787 tmp_dir, num_labels=42, ignore_mismatched_sizes=True788 )789 self.assertIn("the shapes did not match", cl.out)790 791 logits = new_model(**inputs_dict)["logits"]792 self.assertEqual(logits.shape[1], 42)793 794 with CaptureLogger(logger) as cl:795 new_model_without_prefix = FlaxAutoModel.from_pretrained(796 tmp_dir, vocab_size=10, ignore_mismatched_sizes=True797 )798 self.assertIn("the shapes did not match", cl.out)799 input_ids = ids_tensor((2, 8), 10)800 if self.is_encoder_decoder:801 new_model_without_prefix(input_ids, decoder_input_ids=input_ids)802 else:803 new_model_without_prefix(input_ids)804 805 def test_default_params_dtype(self):806 config, _ = self.model_tester.prepare_config_and_inputs_for_common()807 808 for model_class in self.all_model_classes:809 # check if all params are still in float32 when dtype of computation is half-precision810 model = model_class(config, dtype=jnp.float16)811 types = jax.tree_util.tree_map(lambda x: x.dtype, model.params)812 types = flatten_dict(types)813 814 for name, type_ in types.items():815 self.assertEquals(type_, jnp.float32, msg=f"param {name} is not initialized in fp32.")816 817 def test_to_bf16(self):818 config, _ = self.model_tester.prepare_config_and_inputs_for_common()819 820 for model_class in self.all_model_classes:821 model = model_class(config)822 823 # cast all params to bf16824 params = model.to_bf16(model.params)825 types = flatten_dict(jax.tree_util.tree_map(lambda x: x.dtype, params))826 # test if all params are in bf16827 for name, type_ in types.items():828 self.assertEqual(type_, jnp.bfloat16, msg=f"param {name} is not in bf16.")829 830 # test masking831 flat_params = flatten_dict(params)832 key = random.choice(list(flat_params.keys())) # choose a random param833 mask = {path: path != key for path in flat_params} # don't cast the key834 mask = unflatten_dict(mask)835 836 params = model.to_bf16(model.params, mask)837 types = flatten_dict(jax.tree_util.tree_map(lambda x: x.dtype, params))838 # test if all params are in bf16 except key839 for name, type_ in types.items():840 if name == key:841 self.assertEqual(type_, jnp.float32, msg=f"param {name} should be in fp32.")842 else:843 self.assertEqual(type_, jnp.bfloat16, msg=f"param {name} is not in bf16.")844 845 def test_to_fp16(self):846 config, _ = self.model_tester.prepare_config_and_inputs_for_common()847 848 for model_class in self.all_model_classes:849 model = model_class(config)850 851 # cast all params to fp16852 params = model.to_fp16(model.params)853 types = flatten_dict(jax.tree_util.tree_map(lambda x: x.dtype, params))854 # test if all params are in fp16855 for name, type_ in types.items():856 self.assertEqual(type_, jnp.float16, msg=f"param {name} is not in fp16.")857 858 # test masking859 flat_params = flatten_dict(params)860 key = random.choice(list(flat_params.keys())) # choose a random param861 mask = {path: path != key for path in flat_params} # don't cast the key862 mask = unflatten_dict(mask)863 864 params = model.to_fp16(model.params, mask)865 types = flatten_dict(jax.tree_util.tree_map(lambda x: x.dtype, params))866 # test if all params are in fp16 except key867 for name, type_ in types.items():868 if name == key:869 self.assertEqual(type_, jnp.float32, msg=f"param {name} should be in fp32.")870 else:871 self.assertEqual(type_, jnp.float16, msg=f"param {name} is not in fp16.")872 873 def test_to_fp32(self):874 config, _ = self.model_tester.prepare_config_and_inputs_for_common()875 876 for model_class in self.all_model_classes:877 model = model_class(config)878 879 # cast all params to fp16 and back to fp32880 params = model.to_fp16(model.params)881 params = model.to_fp32(params)882 883 # test if all params are in fp32884 types = flatten_dict(jax.tree_util.tree_map(lambda x: x.dtype, params))885 for name, type_ in types.items():886 self.assertEqual(type_, jnp.float32, msg=f"param {name} is not in fp32.")887 888 # test masking889 flat_params = flatten_dict(params)890 key = random.choice(list(flat_params.keys())) # choose a random param891 mask = {path: path != key for path in flat_params} # don't cast the key892 mask = unflatten_dict(mask)893 894 # cast to fp16 and back to fp32 with mask895 params = model.to_fp16(model.params)896 params = model.to_fp32(params, mask)897 898 # test if all params are in fp32 except key899 types = flatten_dict(jax.tree_util.tree_map(lambda x: x.dtype, params))900 for name, type_ in types.items():901 if name == key:902 self.assertEqual(type_, jnp.float16, msg=f"param {name} should be in fp16.")903 else:904 self.assertEqual(type_, jnp.float32, msg=f"param {name} is not in fp32.")905 906 def test_save_load_in_fp16(self):907 config, _ = self.model_tester.prepare_config_and_inputs_for_common()908 909 for model_class in self.all_model_classes:910 model = model_class(config)911 912 # convert weights to fp16 and save913 params = model.to_fp16(model.params)914 with tempfile.TemporaryDirectory() as tmpdirname:915 model.save_pretrained(tmpdirname, params=params)916 917 # load the weights again and check if they are still in fp16918 model = model_class.from_pretrained(tmpdirname)919 types = flatten_dict(jax.tree_util.tree_map(lambda x: x.dtype, model.params))920 for name, type_ in types.items():921 self.assertEqual(type_, jnp.float16, msg=f"param {name} is not in fp16.")922 923 def test_save_load_in_bf16(self):924 config, _ = self.model_tester.prepare_config_and_inputs_for_common()925 926 for model_class in self.all_model_classes:927 model = model_class(config)928 929 # convert weights to bf16 and save930 params = model.to_bf16(model.params)931 with tempfile.TemporaryDirectory() as tmpdirname:932 model.save_pretrained(tmpdirname, params=params)933 934 # load the weights again and check if they are still in fp16935 model = model_class.from_pretrained(tmpdirname)936 types = flatten_dict(jax.tree_util.tree_map(lambda x: x.dtype, model.params))937 for name, type_ in types.items():938 self.assertEqual(type_, jnp.bfloat16, msg=f"param {name} is not in bf16.")939 940 def test_model_main_input_name(self):941 for model_class in self.all_model_classes:942 model_signature = inspect.signature(getattr(model_class, "__call__"))943 # The main input is the name of the argument after `self`944 observed_main_input_name = list(model_signature.parameters.keys())[1]945 self.assertEqual(model_class.main_input_name, observed_main_input_name)946 947 def test_headmasking(self):948 if not self.test_head_masking:949 return950 config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()951 config.return_dict = True952 953 def _prepare_layer_head_mask(i, attention_heads, num_hidden_layers):954 if i == 0:955 return np.concatenate([np.zeros(1, dtype=jnp.int32), np.ones(attention_heads - 1, dtype=jnp.int32)])956 if i == num_hidden_layers - 1:957 return np.concatenate([np.zeros(attention_heads - 1, dtype=jnp.int32), np.ones(1, dtype=jnp.int32)])958 return np.ones(attention_heads, dtype=jnp.int32)959 960 for model_class in self.all_model_classes:961 model = model_class(config)962 963 inputs_dict["output_attentions"] = True964 inputs_dict["output_hidden_states"] = False965 inputs = self._prepare_for_class(inputs_dict, model_class).copy()966 # Prepare head mask967 inputs["head_mask"] = np.stack(968 [969 _prepare_layer_head_mask(i, config.num_attention_heads, config.num_hidden_layers)970 for i in range(config.num_hidden_layers)971 ]972 )973 outputs = model(**inputs)974 975 def _check_attentions_validity(attentions):976 # Remove NaN977 for t in attentions:978 # Check we don't have more than 25% nans (arbitrary)979 self.assertLess(np.isnan(t).sum(), t.size / 4)980 attentions = [np.where(np.isnan(t), 0.0, t) for t in attentions]981 982 self.assertAlmostEqual(attentions[0][..., 0, :, :].sum(), 0.0)983 self.assertNotEqual(attentions[0][..., -1, :, :].sum(), 0.0)984 if len(attentions) > 2: # encoder-decodere models have only 2 layers in each modules985 self.assertNotEqual(attentions[1][..., 0, :, :].sum(), 0.0)986 self.assertAlmostEqual(attentions[-1][..., -2, :, :].sum(), 0.0)987 self.assertNotEqual(attentions[-1][..., -1, :, :].sum(), 0.0)988 989 if model.config.is_encoder_decoder:990 raise NotImplementedError("The test has not been implemented for encoder-decoder models yet.")991 else:992 _check_attentions_validity(outputs.attentions)993 994 def test_no_automatic_init(self):995 config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()996 config.return_dict = True997 998 for model_class in self.all_model_classes:999 model = model_class(config, _do_init=False)1000 1001 # Check that accesing parmas raises an ValueError when _do_init is False1002 with self.assertRaises(ValueError):1003 params = model.params1004 1005 # Check if we params can be properly initialized when calling init_weights1006 params = model.init_weights(model.key, model.input_shape)1007 self.assertIsInstance(params, FrozenDict)1008 # Check if all required parmas are initialized1009 keys = set(flatten_dict(unfreeze(params)).keys())1010 self.assertTrue(all(k in keys for k in model.required_params))1011 # Check if the shapes match1012 flat_params = flatten_dict(unfreeze(params))1013 for k, v in flatten_dict(unfreeze(model.params_shape_tree)).items():1014 self.assertEqual(1015 v.shape,1016 flat_params[k].shape,1017 "Shapes of {} do not match. Expecting {}, got {}.".format(k, v.shape, flat_params[k].shape),1018 )1019 1020 # Check that setting params raises an ValueError when _do_init is False1021 with self.assertRaises(ValueError):1022 model.params = params1023 1024 # Check if we can do a forward pass1025 inputs_dict["output_hidden_states"] = True1026 inputs = self._prepare_for_class(inputs_dict, model_class).copy()1027 model(**inputs, params=params)1028 1029 def test_from_pretrained_with_no_automatic_init(self):1030 config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()1031 config.return_dict = True1032 1033 def _assert_all_params_initialised(model, params):1034 # Check if all required parmas are loaded1035 keys = set(flatten_dict(unfreeze(params)).keys())1036 self.assertTrue(all(k in keys for k in model.required_params))1037 # Check if the shapes match1038 flat_params = flatten_dict(unfreeze(params))1039 for k, v in flatten_dict(unfreeze(model.params_shape_tree)).items():1040 self.assertEqual(1041 v.shape,1042 flat_params[k].shape,1043 "Shapes of {} do not match. Expecting {}, got {}.".format(k, v.shape, flat_params[k].shape),1044 )1045 1046 for model_class in self.all_model_classes:1047 # init the model1048 model = model_class(config)1049 1050 # save the model in the temporary directory1051 # load the saved model with _do_init=False1052 with tempfile.TemporaryDirectory() as tmpdirname:1053 model.save_pretrained(tmpdirname)1054 model, params = model_class.from_pretrained(tmpdirname, _do_init=False)1055 1056 # Check that accesing parmas raises an ValueError when _do_init is False1057 with self.assertRaises(ValueError):1058 params = model.params1059 1060 # Check if all required parmas are loaded1061 _assert_all_params_initialised(model, params)1062 1063 # Check that setting params raises an ValueError when _do_init is False1064 with self.assertRaises(ValueError):1065 model.params = params1066 1067 # Check if init_weights initializes missing keys from from_pretrained1068 flat_params = flatten_dict(unfreeze(params))1069 random_key = random.choice(list(flat_params.keys()))1070 flat_params.pop(random_key)1071 params = freeze(unflatten_dict(flat_params))1072 1073 with tempfile.TemporaryDirectory() as tmpdirname:1074 model.save_pretrained(tmpdirname, params=params)1075 model, params = model_class.from_pretrained(tmpdirname, _do_init=False)1076 1077 params = model.init_weights(model.key, model.input_shape, params=params)1078 # Check if all required parmas are loaded1079 _assert_all_params_initialised(model, params)1080 1081 def test_checkpoint_sharding_from_hub(self):1082 model = FlaxBertModel.from_pretrained("ArthurZ/flax-tiny-random-bert-sharded")1083 # the model above is the same as the model below, just a sharded version.1084 ref_model = FlaxBertModel.from_pretrained("hf-internal-testing/tiny-bert-flax-only")1085 for p1, p2 in zip(flatten_dict(model.params).values(), flatten_dict(ref_model.params).values()):1086 assert np.allclose(np.array(p1), np.array(p2))1087 1088 def test_checkpoint_sharding_local(self):1089 model = FlaxBertModel.from_pretrained("hf-internal-testing/tiny-bert-flax-only")1090 1091 with tempfile.TemporaryDirectory() as tmp_dir:1092 # We use the same folder for various sizes to make sure a new save erases the old checkpoint.1093 for max_size in ["150kB", "150kiB", "200kB", "200kiB"]:1094 model.save_pretrained(tmp_dir, max_shard_size=max_size)1095 1096 # Get each shard file and its size1097 shard_to_size = {}1098 for shard in os.listdir(tmp_dir):1099 if shard.endswith(".msgpack"):1100 shard_file = os.path.join(tmp_dir, shard)1101 shard_to_size[shard_file] = os.path.getsize(shard_file)1102 1103 index_file = os.path.join(tmp_dir, FLAX_WEIGHTS_INDEX_NAME)1104 # Check there is an index but no regular weight file1105 self.assertTrue(os.path.isfile(index_file))1106 self.assertFalse(os.path.isfile(os.path.join(tmp_dir, FLAX_WEIGHTS_NAME)))1107 1108 # Check a file is bigger than max_size only when it has a single weight1109 for shard_file, size in shard_to_size.items():1110 if max_size.endswith("kiB"):1111 max_size_int = int(max_size[:-3]) * 2**101112 else:1113 max_size_int = int(max_size[:-2]) * 10**31114 # Note: pickle adds some junk so the weight of the file can end up being slightly bigger than1115 # the size asked for (since we count parameters)1116 if size >= max_size_int + 50000:1117 with open(shard_file, "rb") as state_f:1118 state_file = from_bytes(FlaxBertModel, state_f.read())1119 self.assertEqual(len(state_file), 1)1120 1121 # Check the index and the shard files found match1122 with open(index_file, "r", encoding="utf-8") as f:1123 index = json.loads(f.read())1124 1125 all_shards = set(index["weight_map"].values())1126 shards_found = {f for f in os.listdir(tmp_dir) if f.endswith(".msgpack")}1127 self.assertSetEqual(all_shards, shards_found)1128 1129 # Finally, check the model can be reloaded1130 new_model = FlaxBertModel.from_pretrained(tmp_dir)1131 for p1, p2 in zip(flatten_dict(model.params).values(), flatten_dict(new_model.params).values()):1132 self.assertTrue(np.allclose(np.array(p1), np.array(p2)))1133 1134 @is_pt_flax_cross_test1135 def test_from_sharded_pt(self):1136 model = FlaxBertModel.from_pretrained("hf-internal-testing/tiny-random-bert-sharded", from_pt=True)1137 ref_model = FlaxBertModel.from_pretrained("hf-internal-testing/tiny-random-bert-fx-only")1138 for key, ref_val in flatten_dict(ref_model.params).items():1139 val = flatten_dict(model.params)[key]1140 assert np.allclose(np.array(val), np.array(ref_val))1141 1142 def test_gradient_checkpointing(self):1143 config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()1144 1145 for model_class in self.all_model_classes:1146 # prepare inputs1147 prepared_inputs_dict = self._prepare_for_class(inputs_dict, model_class)1148 model = model_class(config)1149 remat_model = model_class(config)1150 1151 try:1152 remat_model.enable_gradient_checkpointing()1153 except NotImplementedError:1154 continue1155 1156 outputs = model(**prepared_inputs_dict)1157 remat_outputs = remat_model(**prepared_inputs_dict)1158 1159 # ensure that the dicts of outputs contain the same keys1160 self.assertEqual(outputs.keys(), remat_outputs.keys())1161 1162 outputs = outputs.to_tuple()1163 remat_outputs = remat_outputs.to_tuple()1164 1165 # ensure that the outputs remain precisely equal1166 for output, remat_output in zip(outputs, remat_outputs):1167 self.assertTrue((output == remat_output).all())1168 1169 1170@require_flax1171@is_staging_test1172class FlaxModelPushToHubTester(unittest.TestCase):1173 @classmethod1174 def setUpClass(cls):1175 cls._token = TOKEN1176 HfFolder.save_token(TOKEN)1177 1178 @classmethod1179 def tearDownClass(cls):1180 try:1181 delete_repo(token=cls._token, repo_id="test-model-flax")1182 except HTTPError:1183 pass1184 1185 try:1186 delete_repo(token=cls._token, repo_id="valid_org/test-model-flax-org")1187 except HTTPError:1188 pass1189 1190 def test_push_to_hub(self):1191 config = BertConfig(1192 vocab_size=99, hidden_size=32, num_hidden_layers=5, num_attention_heads=4, intermediate_size=371193 )1194 model = FlaxBertModel(config)1195 model.push_to_hub("test-model-flax", use_auth_token=self._token)1196 1197 new_model = FlaxBertModel.from_pretrained(f"{USER}/test-model-flax")1198 1199 base_params = flatten_dict(unfreeze(model.params))1200 new_params = flatten_dict(unfreeze(new_model.params))