configuration_mt5.py 3.5 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495
  1. # Copyright 2020, The T5 Authors and HuggingFace Inc.
  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 at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # 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 and
  13. # limitations under the License.
  14. """mT5 model configuration"""
  15. from huggingface_hub.dataclasses import strict
  16. from ...configuration_utils import PreTrainedConfig
  17. from ...utils import auto_docstring
  18. @auto_docstring(checkpoint="google/mt5-small")
  19. @strict
  20. class MT5Config(PreTrainedConfig):
  21. r"""
  22. relative_attention_num_buckets (`int`, *optional*, defaults to 32):
  23. The number of buckets to use for each attention layer.
  24. relative_attention_max_distance (`int`, *optional*, defaults to 128):
  25. The maximum distance of the longer sequences for the bucket separation.
  26. feed_forward_proj (`str`, *optional*, defaults to `"gated-gelu"`):
  27. Type of feed forward layer to be used. Should be one of `"relu"` or `"gated-gelu"`.
  28. """
  29. model_type = "mt5"
  30. keys_to_ignore_at_inference = ["past_key_values"]
  31. attribute_map = {
  32. "hidden_size": "d_model",
  33. "num_attention_heads": "num_heads",
  34. "num_hidden_layers": "num_layers",
  35. "head_dim": "d_kv",
  36. }
  37. vocab_size: int = 250112
  38. d_model: int = 512
  39. d_kv: int = 64
  40. d_ff: int = 1024
  41. num_layers: int = 8
  42. num_decoder_layers: int | None = None
  43. num_heads: int = 6
  44. relative_attention_num_buckets: int = 32
  45. relative_attention_max_distance: int = 128
  46. dropout_rate: float | int = 0.1
  47. layer_norm_epsilon: float = 1e-6
  48. initializer_factor: float = 1.0
  49. feed_forward_proj: str = "gated-gelu"
  50. is_encoder_decoder: bool = True
  51. use_cache: bool = True
  52. tie_word_embeddings: bool = True
  53. bos_token_id: int | None = None
  54. pad_token_id: int | None = 0
  55. eos_token_id: int | list[int] | None = 1
  56. decoder_start_token_id: int | None = 0
  57. classifier_dropout: float | int = 0.0
  58. is_decoder: bool = False
  59. def __post_init__(self, **kwargs):
  60. self.num_decoder_layers = (
  61. self.num_decoder_layers if self.num_decoder_layers is not None else self.num_layers
  62. ) # default = symmetry
  63. act_info = self.feed_forward_proj.split("-")
  64. self.dense_act_fn = act_info[-1]
  65. self.is_gated_act = act_info[0] == "gated"
  66. if self.feed_forward_proj == "gated-gelu":
  67. self.dense_act_fn = "gelu_new"
  68. # Force because official weights have False serialized, but we have to tie always
  69. kwargs.pop("tie_word_embeddings", None)
  70. self.tie_word_embeddings = True
  71. super().__post_init__(**kwargs)
  72. def validate_architecture(self):
  73. """Part of `@strict`-powered validation. Validates the architecture of the config."""
  74. act_info = self.feed_forward_proj.split("-")
  75. if len(act_info) > 1 and act_info[0] != "gated" or len(act_info) > 2:
  76. raise ValueError(
  77. f"`feed_forward_proj`: {self.feed_forward_proj} is not a valid activation function of the dense layer. "
  78. "Please make sure `feed_forward_proj` is of the format `gated-{ACT_FN}` or `{ACT_FN}`, e.g. "
  79. "'gated-gelu' or 'relu'"
  80. )
  81. __all__ = ["MT5Config"]