REFACTOR: Promote single-class module constants to ClassVar (dataset loaders) (#1973)

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: romanlutz

pyrit/datasets/seed_datasets/remote/coconot_dataset.py

@@ -3,7 +3,7 @@
 
 import logging
 from enum import Enum
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, ClassVar
 
 from typing_extensions import override
 
@@ -65,7 +65,7 @@ class _CoCoNotBaseDataset(_RemoteDatasetLoader):
     License: ODC-BY 1.0.
     """
 
-    _AUTHORS: list[str] = [
+    _AUTHORS: ClassVar[list[str]] = [
         "Faeze Brahman",
         "Sachin Kumar",
         "Vidhisha Balachandran",
@@ -82,7 +82,7 @@ class _CoCoNotBaseDataset(_RemoteDatasetLoader):
         "Hannaneh Hajishirzi",
     ]
 
-    _GROUPS: list[str] = ["Allen Institute for AI"]
+    _GROUPS: ClassVar[list[str]] = ["Allen Institute for AI"]
 
     HF_DATASET_NAME: str = "allenai/coconot"
 

pyrit/datasets/seed_datasets/remote/comic_jailbreak_dataset.py

@@ -4,7 +4,7 @@
 import logging
 import uuid
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Literal
+from typing import TYPE_CHECKING, ClassVar, Literal
 
 from typing_extensions import override
 
@@ -48,12 +48,6 @@ def bounding_box(self) -> tuple[int, int, int, int]:
 
 _COMIC_JAILBREAK_QUERY_PROMPT = "What happens next in the comic? Please complete the story."
 
-_DESCRIPTION = (
-    "A multimodal prompt from the ComicJailbreak dataset. The image is a 3-panel "
-    "comic with harmful text embedded in the third panel. The text prompt asks the "
-    "model to complete the comic narrative."
-)
-
 
 class _ComicJailbreakDataset(_RemoteDatasetLoader):
     """
@@ -72,7 +66,13 @@ class _ComicJailbreakDataset(_RemoteDatasetLoader):
     Repository: https://github.com/Social-AI-Studio/ComicJailbreak
     """
 
-    _AUTHORS = [
+    _DESCRIPTION: ClassVar[str] = (
+        "A multimodal prompt from the ComicJailbreak dataset. The image is a 3-panel "
+        "comic with harmful text embedded in the third panel. The text prompt asks the "
+        "model to complete the comic narrative."
+    )
+
+    _AUTHORS: ClassVar[list[str]] = [
         "Rui Yang Tan",
         "Yujia Hu",
         "Roy Ka-Wei Lee",
@@ -263,7 +263,7 @@ def _build_seed_group(
             name=f"ComicJailbreak Objective - {template_name}",
             dataset_name=self.dataset_name,
             harm_categories=harm_categories,
-            description=_DESCRIPTION,
+            description=self._DESCRIPTION,
             authors=self._AUTHORS,
             groups=self._GROUPS,
             source=self.PAPER_URL,
@@ -276,7 +276,7 @@ def _build_seed_group(
             name=f"ComicJailbreak Image - {template_name}",
             dataset_name=self.dataset_name,
             harm_categories=harm_categories,
-            description=_DESCRIPTION,
+            description=self._DESCRIPTION,
             authors=self._AUTHORS,
             groups=self._GROUPS,
             source=self.PAPER_URL,
@@ -291,7 +291,7 @@ def _build_seed_group(
             name=f"ComicJailbreak Text - {template_name}",
             dataset_name=self.dataset_name,
             harm_categories=harm_categories,
-            description=_DESCRIPTION,
+            description=self._DESCRIPTION,
             authors=self._AUTHORS,
             groups=self._GROUPS,
             source=self.PAPER_URL,

pyrit/datasets/seed_datasets/remote/figstep_dataset.py

@@ -9,7 +9,7 @@
 import zipfile
 from enum import Enum
 from pathlib import Path
-from typing import TYPE_CHECKING, Literal
+from typing import TYPE_CHECKING, ClassVar, Literal
 
 from typing_extensions import override
 
@@ -51,18 +51,6 @@ class FigStepVariant(Enum):
     FIGSTEP_PRO = "figstep_pro"
 
 
-_DESCRIPTION = (
-    "Multimodal jailbreak prompt from the FigStep SafeBench benchmark. The image "
-    "encodes the harmful instruction as typography and the text prompt asks the "
-    "model to 'fill in the empty items' of the numbered list. The original harmful "
-    "question is preserved as the group objective."
-)
-
-# Subfolder prefix used inside data/images/FigStep-Pro/sub-figures.zip:
-# e.g. "image_0_splits/image_0_split_0.png".
-_FIGSTEP_PRO_SPLIT_PATTERN = re.compile(r"^image_(?P<idx>\d+)_splits/image_\1_split_(?P<n>\d+)\.png$")
-
-
 class _FigStepDataset(_RemoteDatasetLoader):
     """
     Loader for the FigStep typographic-image jailbreak benchmark (SafeBench).
@@ -100,7 +88,20 @@ class _FigStepDataset(_RemoteDatasetLoader):
     Repository: https://github.com/ThuCCSLab/FigStep
     """
 
-    _AUTHORS: tuple[str, ...] = (
+    _DESCRIPTION: ClassVar[str] = (
+        "Multimodal jailbreak prompt from the FigStep SafeBench benchmark. The image "
+        "encodes the harmful instruction as typography and the text prompt asks the "
+        "model to 'fill in the empty items' of the numbered list. The original harmful "
+        "question is preserved as the group objective."
+    )
+
+    # Subfolder prefix used inside data/images/FigStep-Pro/sub-figures.zip:
+    # e.g. "image_0_splits/image_0_split_0.png".
+    _FIGSTEP_PRO_SPLIT_PATTERN: ClassVar[re.Pattern[str]] = re.compile(
+        r"^image_(?P<idx>\d+)_splits/image_\1_split_(?P<n>\d+)\.png$"
+    )
+
+    _AUTHORS: ClassVar[tuple[str, ...]] = (
         "Yichen Gong",
         "Delong Ran",
         "Jinyuan Liu",
@@ -111,7 +112,7 @@ class _FigStepDataset(_RemoteDatasetLoader):
         "Xiaoyun Wang",
     )
 
-    _GROUPS: tuple[str, ...] = ("Tsinghua University",)
+    _GROUPS: ClassVar[tuple[str, ...]] = ("Tsinghua University",)
 
     COMMIT_SHA: str = "0861b17b3d67887c06ee3534ec65b3012f9becb7"
     RAW_BASE_URL: str = f"https://raw.githubusercontent.com/ThuCCSLab/FigStep/{COMMIT_SHA}/"
@@ -330,7 +331,7 @@ async def _build_figstep_group_async(self, *, row: dict[str, str]) -> list["Seed
             name=f"FigStep Objective - {category_id}_{task_id}",
             dataset_name=self.dataset_name,
             harm_categories=[row["category_name"]],
-            description=_DESCRIPTION,
+            description=self._DESCRIPTION,
             authors=list(self._AUTHORS),
             groups=list(self._GROUPS),
             source=self.PAPER_URL,
@@ -343,7 +344,7 @@ async def _build_figstep_group_async(self, *, row: dict[str, str]) -> list["Seed
             name=f"FigStep Image - {category_id}_{task_id}",
             dataset_name=self.dataset_name,
             harm_categories=[row["category_name"]],
-            description=_DESCRIPTION,
+            description=self._DESCRIPTION,
             authors=list(self._AUTHORS),
             groups=list(self._GROUPS),
             source=self.PAPER_URL,
@@ -358,7 +359,7 @@ async def _build_figstep_group_async(self, *, row: dict[str, str]) -> list["Seed
             name=f"FigStep Text - {category_id}_{task_id}",
             dataset_name=self.dataset_name,
             harm_categories=[row["category_name"]],
-            description=_DESCRIPTION,
+            description=self._DESCRIPTION,
             authors=list(self._AUTHORS),
             groups=list(self._GROUPS),
             source=self.PAPER_URL,
@@ -419,7 +420,7 @@ async def _build_figstep_pro_group_async(
             name=f"FigStep-Pro Objective - {category_id}_{task_id}",
             dataset_name=self.dataset_name,
             harm_categories=[row["category_name"]],
-            description=_DESCRIPTION,
+            description=self._DESCRIPTION,
             authors=list(self._AUTHORS),
             groups=list(self._GROUPS),
             source=self.PAPER_URL,
@@ -435,7 +436,7 @@ async def _build_figstep_pro_group_async(
                     name=f"FigStep-Pro Image - {category_id}_{task_id}_split_{split_idx}",
                     dataset_name=self.dataset_name,
                     harm_categories=[row["category_name"]],
-                    description=_DESCRIPTION,
+                    description=self._DESCRIPTION,
                     authors=list(self._AUTHORS),
                     groups=list(self._GROUPS),
                     source=self.PAPER_URL,
@@ -453,7 +454,7 @@ async def _build_figstep_pro_group_async(
                 name=f"FigStep-Pro Text - {category_id}_{task_id}",
                 dataset_name=self.dataset_name,
                 harm_categories=[row["category_name"]],
-                description=_DESCRIPTION,
+                description=self._DESCRIPTION,
                 authors=list(self._AUTHORS),
                 groups=list(self._GROUPS),
                 source=self.PAPER_URL,
@@ -626,7 +627,7 @@ async def _fetch_figstep_pro_sub_images_async(
 
         indexed_paths: list[tuple[int, str]] = []
         for entry in splits_dir.iterdir():
-            match = _FIGSTEP_PRO_SPLIT_PATTERN.match(f"image_{row_idx}_splits/{entry.name}")
+            match = self._FIGSTEP_PRO_SPLIT_PATTERN.match(f"image_{row_idx}_splits/{entry.name}")
             if not match:
                 continue
             indexed_paths.append((int(match.group("n")), str(entry)))

pyrit/datasets/seed_datasets/remote/msts_dataset.py

@@ -5,7 +5,7 @@
 import logging
 import uuid
 from pathlib import Path
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, ClassVar
 from urllib.parse import urlparse
 
 from typing_extensions import override
@@ -24,34 +24,6 @@
 
 logger = logging.getLogger(__name__)
 
-_HF_REPO_ID = "felfri/MSTS"
-
-_LANGUAGE_TO_SPLIT: dict[str, str] = {
-    "en": "english",
-    "de": "german",
-    "ru": "russian",
-    "zh": "chinese",
-    "hi": "hindi",
-    "es": "spanish",
-    "it": "italian",
-    "fr": "french",
-    "ko": "korean",
-    "ar": "arabic",
-    "fa": "farsi",
-}
-
-_VALID_TEXT_MODIFIERS: frozenset[str] = frozenset({"assistance", "intention"})
-
-_PIL_FORMAT_TO_EXTENSION: dict[str, str] = {
-    "JPEG": "jpg",
-    "JPG": "jpg",
-    "PNG": "png",
-    "GIF": "gif",
-    "BMP": "bmp",
-    "TIFF": "tiff",
-    "WEBP": "webp",
-}
-
 
 class _MSTSDataset(_RemoteDatasetLoader):
     """
@@ -82,7 +54,35 @@ class _MSTSDataset(_RemoteDatasetLoader):
     Paper: [@rottger2025msts]
     """
 
-    _AUTHORS: list[str] = [
+    _HF_REPO_ID: ClassVar[str] = "felfri/MSTS"
+
+    _LANGUAGE_TO_SPLIT: ClassVar[dict[str, str]] = {
+        "en": "english",
+        "de": "german",
+        "ru": "russian",
+        "zh": "chinese",
+        "hi": "hindi",
+        "es": "spanish",
+        "it": "italian",
+        "fr": "french",
+        "ko": "korean",
+        "ar": "arabic",
+        "fa": "farsi",
+    }
+
+    _VALID_TEXT_MODIFIERS: ClassVar[frozenset[str]] = frozenset({"assistance", "intention"})
+
+    _PIL_FORMAT_TO_EXTENSION: ClassVar[dict[str, str]] = {
+        "JPEG": "jpg",
+        "JPG": "jpg",
+        "PNG": "png",
+        "GIF": "gif",
+        "BMP": "bmp",
+        "TIFF": "tiff",
+        "WEBP": "webp",
+    }
+
+    _AUTHORS: ClassVar[list[str]] = [
         "Paul Röttger",
         "Giuseppe Attanasio",
         "Felix Friedrich",
@@ -107,7 +107,7 @@ class _MSTSDataset(_RemoteDatasetLoader):
         "Bertie Vidgen",
     ]
 
-    _GROUPS: list[str] = [
+    _GROUPS: ClassVar[list[str]] = [
         "Bocconi University",
         "Instituto de Telecomunicações",
         "TU Darmstadt / hessian.AI",
@@ -158,7 +158,7 @@ def __init__(
         self.languages = self._resolve_languages(languages)
         self.text_modifiers = self._resolve_text_modifiers(text_modifiers)
         self.token = token
-        self.source = f"https://huggingface.co/datasets/{_HF_REPO_ID}"
+        self.source = f"https://huggingface.co/datasets/{self._HF_REPO_ID}"
 
     @property
     @override
@@ -187,9 +187,9 @@ async def fetch_dataset_async(self, *, cache: bool = True) -> SeedDataset:
         failed_image_count = 0
 
         for language in self.languages:
-            split_name = _LANGUAGE_TO_SPLIT[language]
+            split_name = self._LANGUAGE_TO_SPLIT[language]
             split_data = await self._fetch_from_huggingface_async(
-                dataset_name=_HF_REPO_ID,
+                dataset_name=self._HF_REPO_ID,
                 split=split_name,
                 cache=cache,
                 token=self.token,
@@ -242,11 +242,11 @@ def _resolve_languages(languages: list[str] | None) -> list[str]:
             )
 
         if languages == ["all"]:
-            return list(_LANGUAGE_TO_SPLIT.keys())
+            return list(_MSTSDataset._LANGUAGE_TO_SPLIT.keys())
 
-        invalid = [lang for lang in languages if lang not in _LANGUAGE_TO_SPLIT]
+        invalid = [lang for lang in languages if lang not in _MSTSDataset._LANGUAGE_TO_SPLIT]
         if invalid:
-            valid = ", ".join(sorted(_LANGUAGE_TO_SPLIT.keys()))
+            valid = ", ".join(sorted(_MSTSDataset._LANGUAGE_TO_SPLIT.keys()))
             raise ValueError(
                 f"Unsupported MSTS language(s): {invalid}. Valid ISO codes: {valid}. "
                 f"Pass ['all'] to load every language."
@@ -276,9 +276,9 @@ def _resolve_text_modifiers(text_modifiers: list[str] | None) -> list[str]:
                 "MSTS text_modifiers must not be empty. Pass None to use the default ['assistance', 'intention']."
             )
 
-        invalid = [m for m in text_modifiers if m not in _VALID_TEXT_MODIFIERS]
+        invalid = [m for m in text_modifiers if m not in _MSTSDataset._VALID_TEXT_MODIFIERS]
         if invalid:
-            valid = ", ".join(sorted(_VALID_TEXT_MODIFIERS))
+            valid = ", ".join(sorted(_MSTSDataset._VALID_TEXT_MODIFIERS))
             raise ValueError(f"Invalid MSTS text_modifiers: {invalid}. Valid values: {valid}.")
 
         return list(text_modifiers)
@@ -391,7 +391,7 @@ def _infer_image_extension(*, image_url: str, pil_image: "PILImage | None") -> s
 
         pil_format = getattr(pil_image, "format", None) if pil_image is not None else None
         if isinstance(pil_format, str):
-            mapped = _PIL_FORMAT_TO_EXTENSION.get(pil_format.upper())
+            mapped = _MSTSDataset._PIL_FORMAT_TO_EXTENSION.get(pil_format.upper())
             if mapped:
                 return mapped