configuration_umt5.py 3.4 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091
  1. # Copyright 2023, 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. """UMT5 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/umt5-small")
  19. @strict
  20. class UMT5Config(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 = "umt5"
  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. pad_token_id: int | None = 0
  53. eos_token_id: int | list[int] | None = 1
  54. decoder_start_token_id: int | None = 0
  55. classifier_dropout: float | int = 0.0
  56. is_decoder: bool = False
  57. def __post_init__(self, **kwargs):
  58. self.num_decoder_layers = (
  59. self.num_decoder_layers if self.num_decoder_layers is not None else self.num_layers
  60. ) # default = symmetry
  61. act_info = self.feed_forward_proj.split("-")
  62. self.dense_act_fn = act_info[-1]
  63. self.is_gated_act = act_info[0] == "gated"
  64. if self.feed_forward_proj == "gated-gelu":
  65. self.dense_act_fn = "gelu_new"
  66. kwargs.pop("tie_word_embeddings", None)
  67. self.tie_word_embeddings = True # force it for T5 family
  68. super().__post_init__(**kwargs)
  69. def validate_architecture(self):
  70. """Part of `@strict`-powered validation. Validates the architecture of the config."""
  71. act_info = self.feed_forward_proj.split("-")
  72. if len(act_info) > 1 and act_info[0] != "gated" or len(act_info) > 2:
  73. raise ValueError(
  74. f"`feed_forward_proj`: {self.feed_forward_proj} is not a valid activation function of the dense layer. "
  75. "Please make sure `feed_forward_proj` is of the format `gated-{ACT_FN}` or `{ACT_FN}`, e.g. "
  76. "'gated-gelu' or 'relu'"
  77. )
  78. __all__ = ["UMT5Config"]