configuration_shieldgemma2.py 4.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118
  1. # coding=utf-8
  2. # Copyright 2025 Google Inc. HuggingFace Inc. team. All rights reserved.
  3. #
  4. #
  5. # Licensed under the Apache License, Version 2.0 (the "License");
  6. # you may not use this file except in compliance with the License.
  7. # You may obtain a copy of the License at
  8. #
  9. # http://www.apache.org/licenses/LICENSE-2.0
  10. #
  11. # Unless required by applicable law or agreed to in writing, software
  12. # distributed under the License is distributed on an "AS IS" BASIS,
  13. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  14. # See the License for the specific language governing permissions and
  15. # limitations under the License.
  16. from ...configuration_utils import PretrainedConfig
  17. from ...utils import logging
  18. from ..auto import CONFIG_MAPPING, AutoConfig
  19. logger = logging.get_logger(__name__)
  20. class ShieldGemma2Config(PretrainedConfig):
  21. r"""
  22. This is the configuration class to store the configuration of a [`ShieldGemma2ForImageClassification`]. It is used to instantiate an
  23. ShieldGemma2ForImageClassification according to the specified arguments, defining the model architecture. Instantiating a configuration
  24. with the defaults will yield a similar configuration to that of the shieldgemma-2-4b-it.
  25. e.g. [google/gemma-3-4b](https://huggingface.co/google/gemma-3-4b)
  26. Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
  27. documentation from [`PretrainedConfig`] for more information.
  28. Args:
  29. text_config (`Union[ShieldGemma2TextConfig, dict]`, *optional*):
  30. The config object of the text backbone.
  31. vision_config (`Union[AutoConfig, dict]`, *optional*):
  32. Custom vision config or dict.
  33. mm_tokens_per_image (`int`, *optional*, defaults to 256):
  34. The number of tokens per image embedding.
  35. boi_token_index (`int`, *optional*, defaults to 255999):
  36. The begin-of-image token index to wrap the image prompt.
  37. eoi_token_index (`int`, *optional*, defaults to 256000):
  38. The end-of-image token index to wrap the image prompt.
  39. image_token_index (`int`, *optional*, defaults to 262144):
  40. The image token index to encode the image prompt.
  41. initializer_range (`float`, *optional*, defaults to 0.02):
  42. The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
  43. Example:
  44. ```python
  45. >>> from transformers import ShieldGemma2ForConditionalGeneration, ShieldGemma2Config, SiglipVisionConfig, ShieldGemma2TextConfig
  46. >>> # Initializing a Siglip-like vision config
  47. >>> vision_config = SiglipVisionConfig()
  48. >>> # Initializing a ShieldGemma2 Text config
  49. >>> text_config = ShieldGemma2TextConfig()
  50. >>> # Initializing a ShieldGemma2 gemma-3-4b style configuration
  51. >>> configuration = ShieldGemma2Config(vision_config, text_config)
  52. >>> # Initializing a model from the gemma-3-4b style configuration
  53. >>> model = ShieldGemma2TextConfig(configuration)
  54. >>> # Accessing the model configuration
  55. >>> configuration = model.config
  56. ```"""
  57. model_type = "shieldgemma2"
  58. attribute_map = {
  59. "image_token_id": "image_token_index",
  60. "boi_token_id": "boi_token_index",
  61. "eoi_token_id": "eoi_token_index",
  62. }
  63. sub_configs = {"text_config": AutoConfig, "vision_config": AutoConfig}
  64. def __init__(
  65. self,
  66. text_config=None,
  67. vision_config=None,
  68. mm_tokens_per_image: int = 256,
  69. boi_token_index: int = 255_999,
  70. eoi_token_index: int = 256_000,
  71. image_token_index: int = 262_144,
  72. initializer_range: float = 0.02,
  73. **kwargs,
  74. ):
  75. if isinstance(vision_config, dict):
  76. vision_config["model_type"] = vision_config.get("model_type", "siglip_vision_model")
  77. vision_config = CONFIG_MAPPING[vision_config["model_type"]](**vision_config)
  78. elif vision_config is None:
  79. vision_config = CONFIG_MAPPING["siglip_vision_model"]()
  80. self.vision_config = vision_config
  81. if isinstance(text_config, dict):
  82. text_config["model_type"] = text_config.get("model_type", "gemma3_text")
  83. text_config = CONFIG_MAPPING[text_config["model_type"]](**text_config)
  84. elif text_config is None:
  85. text_config = CONFIG_MAPPING["gemma3_text"]()
  86. self.text_config = text_config
  87. self.vision_config = vision_config
  88. self.mm_tokens_per_image = mm_tokens_per_image
  89. self.boi_token_index = boi_token_index
  90. self.eoi_token_index = eoi_token_index
  91. self.image_token_index = image_token_index
  92. self.initializer_range = initializer_range
  93. super().__init__(**kwargs)
  94. __all__ = ["ShieldGemma2Config"]