configuration_shieldgemma2.py 3.6 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192
  1. # Copyright 2025 Google Inc. HuggingFace Inc. team. All rights reserved.
  2. #
  3. #
  4. # Licensed under the Apache License, Version 2.0 (the "License");
  5. # you may not use this file except in compliance with the License.
  6. # You may obtain a copy of the License at
  7. #
  8. # http://www.apache.org/licenses/LICENSE-2.0
  9. #
  10. # Unless required by applicable law or agreed to in writing, software
  11. # distributed under the License is distributed on an "AS IS" BASIS,
  12. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13. # See the License for the specific language governing permissions and
  14. # limitations under the License.
  15. from huggingface_hub.dataclasses import strict
  16. from ...configuration_utils import PreTrainedConfig
  17. from ...utils import auto_docstring
  18. from ..auto import CONFIG_MAPPING, AutoConfig
  19. @auto_docstring(checkpoint="google/shieldgemma-2-4b-it")
  20. @strict
  21. class ShieldGemma2Config(PreTrainedConfig):
  22. r"""
  23. tie_word_embeddings (`bool`, *optional*):
  24. Whether to tie the word embeddings. Defaults to the value of `text_config.tie_word_embeddings` if not set.
  25. mm_tokens_per_image (`int`, *optional*, defaults to 256):
  26. The number of tokens per image embedding.
  27. boi_token_index (`int`, *optional*, defaults to 255999):
  28. The begin-of-image token index to wrap the image prompt.
  29. eoi_token_index (`int`, *optional*, defaults to 256000):
  30. The end-of-image token index to wrap the image prompt.
  31. Example:
  32. ```python
  33. >>> from transformers import ShieldGemma2ForConditionalGeneration, ShieldGemma2Config, SiglipVisionConfig, ShieldGemma2TextConfig
  34. >>> # Initializing a Siglip-like vision config
  35. >>> vision_config = SiglipVisionConfig()
  36. >>> # Initializing a ShieldGemma2 Text config
  37. >>> text_config = ShieldGemma2TextConfig()
  38. >>> # Initializing a ShieldGemma2 gemma-3-4b style configuration
  39. >>> configuration = ShieldGemma2Config(vision_config, text_config)
  40. >>> # Initializing a model from the gemma-3-4b style configuration
  41. >>> model = ShieldGemma2TextConfig(configuration)
  42. >>> # Accessing the model configuration
  43. >>> configuration = model.config
  44. ```"""
  45. model_type = "shieldgemma2"
  46. attribute_map = {
  47. "image_token_id": "image_token_index",
  48. "boi_token_id": "boi_token_index",
  49. "eoi_token_id": "eoi_token_index",
  50. }
  51. sub_configs = {"text_config": AutoConfig, "vision_config": AutoConfig}
  52. text_config: dict | PreTrainedConfig | None = None
  53. vision_config: dict | PreTrainedConfig | None = None
  54. mm_tokens_per_image: int = 256
  55. boi_token_index: int = 255_999
  56. eoi_token_index: int = 256_000
  57. image_token_index: int = 262_144
  58. initializer_range: float = 0.02
  59. def __post_init__(self, **kwargs):
  60. if isinstance(self.vision_config, dict):
  61. self.vision_config["model_type"] = self.vision_config.get("model_type", "siglip_vision_model")
  62. self.vision_config = CONFIG_MAPPING[self.vision_config["model_type"]](**self.vision_config)
  63. elif self.vision_config is None:
  64. self.vision_config = CONFIG_MAPPING["siglip_vision_model"]()
  65. if isinstance(self.text_config, dict):
  66. self.text_config["model_type"] = self.text_config.get("model_type", "gemma3_text")
  67. self.text_config = CONFIG_MAPPING[self.text_config["model_type"]](**self.text_config)
  68. elif self.text_config is None:
  69. self.text_config = CONFIG_MAPPING["gemma3_text"]()
  70. if kwargs.get("tie_word_embeddings") is None:
  71. self.tie_word_embeddings = getattr(self.text_config, "tie_word_embeddings", True)
  72. super().__post_init__(**kwargs)
  73. __all__ = ["ShieldGemma2Config"]