Fixing empty messages bug reported by ghunkins

Author: RyanMullins

docs/source/en/model_doc/shieldgemma2.md

@@ -85,7 +85,16 @@ output = model(**inputs)
 print(output.probabilities)
 ```
 
-## ShieldGemmaForImageClassification
 
-[[autodoc]] Shieldgemma2ForTokenClassification
+## ShieldGemma2Processor
+
+[[autodoc]] ShieldGemma2Processor
+
+## ShieldGemma2Config
+
+[[autodoc]] ShieldGemma2Config
+
+## ShieldGemma2ForImageClassification
+
+[[autodoc]] ShieldGemma2ForImageClassification
     - forward

src/transformers/models/shieldgemma2/modeling_shieldgemma2.py

@@ -212,7 +212,7 @@ def forward(
             output_hidden_states=output_hidden_states,
             return_dict=return_dict,
             logits_to_keep=logits_to_keep,
-            **lm_kwargs
+            **lm_kwargs,
         )
         logits = outputs.logits
         selected_logits = logits[:, -1, [self.yes_token_index, self.no_token_index]]

src/transformers/models/shieldgemma2/processing_shieldgemma2.py

@@ -25,6 +25,26 @@
 
 logger = logging.get_logger(__name__)
 
+DEFAULT_SHIELDGEMMA2_POLICIES: Mapping[str, str] = {
+    "dangerous": (
+        "No Dangerous Content: The image shall not contain content that"
+        " facilitates or encourages activities that could cause real-world harm"
+        " (e.g., building firearms and explosive devices, promotion of"
+        " terrorism, instructions for suicide)."
+    ),
+    "sexual": (
+        "No Sexually Explicit content: The image shall not contain content that"
+        " depicts explicit or graphic sexual acts (e.g., pornography, erotic"
+        " nudity, depictions of rape or sexual assault)."
+    ),
+    "violence": (
+        "No Violence/Gore content: The image shall not contain content that"
+        " depicts shocking, sensational, or gratuitous violence (e.g.,"
+        " excessive blood and gore, gratuitous violence against animals,"
+        " extreme injury or moment of death)."
+    ),
+}
+
 
 class ShieldGemma2ProcessorKwargs(Gemma3ProcessorKwargs, total=False):
     policies: Optional[Sequence[str]]
@@ -40,15 +60,8 @@ class ShieldGemma2ProcessorKwargs(Gemma3ProcessorKwargs, total=False):
 
 
 class ShieldGemma2Processor(Gemma3Processor):
-
     def __init__(
-        self,
-        image_processor,
-        tokenizer,
-        chat_template = None,
-        image_seq_length = 256,
-        policy_definitions = None,
-        **kwargs
+        self, image_processor, tokenizer, chat_template=None, image_seq_length=256, policy_definitions=None, **kwargs
     ):
         """A processor for the ShieldGemma 2 model.
 
@@ -65,10 +78,10 @@ def __init__(
                 the base policies ShieldGemma was trained on.
         """
         super().__init__(image_processor, tokenizer, chat_template, image_seq_length, **kwargs)
-        if policy_definitions:
-            self.policy_definitions = policy_definitions
+        if policy_definitions is None:
+            self.policy_definitions = DEFAULT_SHIELDGEMMA2_POLICIES
         else:
-            self.policy_definitions = {}
+            self.policy_definitions = policy_definitions
 
     def __call__(
         self,
@@ -129,7 +142,6 @@ def __call__(
             text_kwargs["padding"] = kwargs.pop("padding", True)
             text_kwargs["padding_side"] = kwargs.pop("padding_side", "left")
 
-
         policy_definitions: Mapping[str, str] = {
             **self.policy_definitions,
             **kwargs.get("custom_policies", {}),
@@ -138,6 +150,7 @@ def __call__(
         if (policies := kwargs.get("policies")) is None:
             policies = list(policy_definitions.keys())
 
+        # TODO(ryanmullins): Support images from PIL or URLs.
         messages = []
         expanded_images = []
         for img in images:

tests/models/shieldgemma2/test_modeling_shieldgemma2.py

@@ -14,11 +14,11 @@
 # limitations under the License.
 """Testing suite for the PyTorch Gemma3 model."""
 
-from io import BytesIO
 import unittest
+from io import BytesIO
 
-from PIL import Image
 import requests
+from PIL import Image
 
 from transformers import is_torch_available
 from transformers.testing_utils import (
@@ -31,14 +31,14 @@
 
 if is_torch_available():
     import torch
+
     from transformers import ShieldGemma2ForImageClassification, ShieldGemma2Processor
 
 
 @slow
 @require_torch_gpu
 # @require_read_token
 class ShieldGemma2IntegrationTest(unittest.TestCase):
-
     def tearDown(self):
         cleanup(torch_device, gc_collect=True)
 

tests/models/shieldgemma2/test_processing_shieldgemma2.py

@@ -12,14 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from collections.abc import Mapping
 import json
 import os
 import shutil
 import tempfile
 import unittest
+from collections.abc import Mapping
 
 from parameterized import parameterized
+
 from transformers import GemmaTokenizer, ShieldGemma2Processor
 from transformers.testing_utils import get_tests_dir, require_vision
 from transformers.utils import is_vision_available
@@ -95,29 +96,6 @@ def prepare_processor_dict(self):
             "policy_definitions": _SHIELDGEMMA2_POLICIES,
         }
 
-    def test_unstructured_kwargs(self):
-        if "image_processor" not in self.processor_class.attributes:
-            self.skipTest(f"image_processor attribute not present in {self.processor_class}")
-        processor_components = self.prepare_components()
-        processor_kwargs = self.prepare_processor_dict()
-        processor = self.processor_class(**processor_components, **processor_kwargs)
-        self.skip_processor_without_typed_kwargs(processor)
-
-        input_str = self.prepare_text_inputs()
-        image_input = self.prepare_image_inputs()
-        inputs = processor(
-            text=input_str,
-            images=image_input,
-            return_tensors="pt",
-            do_rescale=True,
-            rescale_factor=-1,
-            padding="max_length",
-            max_length=674,
-        )
-
-        self.assertLessEqual(inputs[self.images_input_name][0][0].mean(), 0)
-        self.assertEqual(inputs[self.text_input_name].shape[-1], 674)
-
     def test_policy_definitions_saved_in_config(self):
         processor_config_path = os.path.join(self.tmpdirname, "processor_config.json")
 
@@ -128,11 +106,13 @@ def test_policy_definitions_saved_in_config(self):
         self.assertIn("policy_definitions", json_dict)
         self.assertIs(len(json_dict["policy_definitions"]), 3)
 
-    @parameterized.expand([
-        ("all_policies", None, 3),
-        ("selected_policies", ["dangerous", "violence"], 2),
-        ("single_policy", ["sexual"], 1),
-    ])
+    @parameterized.expand(
+        [
+            ("all_policies", None, 3),
+            ("selected_policies", ["dangerous", "violence"], 2),
+            ("single_policy", ["sexual"], 1),
+        ]
+    )
     def test_with_default_policies(self, name, policies, expected_batch_size):
         processor = self.get_processor()
 
@@ -144,14 +124,16 @@ def test_with_default_policies(self, name, policies, expected_batch_size):
         self.assertEqual(len(processed_inputs[self.text_input_name]), expected_batch_size)
         self.assertEqual(len(processed_inputs[self.images_input_name]), expected_batch_size)
 
-    @parameterized.expand([
-        ("all_policies", None, 6),
-        ("selected_policies_from_both", ["cbrne", "dangerous", "specialized_advice", "violence"], 4),
-        ("selected_policies_from_custom", ["cbrne", "specialized_advice"], 2),
-        ("selected_policies_from_default", ["dangerous", "violence"], 2),
-        ("single_policy_from_custom", ["ip"], 1),
-        ("single_policy_from_default", ["sexual"], 1),
-    ])
+    @parameterized.expand(
+        [
+            ("all_policies", None, 6),
+            ("selected_policies_from_both", ["cbrne", "dangerous", "specialized_advice", "violence"], 4),
+            ("selected_policies_from_custom", ["cbrne", "specialized_advice"], 2),
+            ("selected_policies_from_default", ["dangerous", "violence"], 2),
+            ("single_policy_from_custom", ["ip"], 1),
+            ("single_policy_from_default", ["sexual"], 1),
+        ]
+    )
     def test_with_custom_policies(self, name, policies, expected_batch_size):
         processor = self.get_processor()