scripts/adapters/b2_mmsafetybench_adapter.py

"""
MM-SafetyBench Dataset Adapter for B2 (Image Poisoning) Schema

MM-SafetyBench: Multimodal safety benchmark (PKU-Alignment)
License: CC-BY-NC-4.0 (RESEARCH ONLY - DO NOT REDISTRIBUTE)
HuggingFace: https://huggingface.co/datasets/PKU-Alignment/MM-SafetyBench

This adapter maps MM-SafetyBench multimodal safety prompts into the benchmark's
B2 schema. Many MM-SafetyBench prompts reference text embedded in the image
("TYPO"/"SD_TYPO" splits), which aligns with the benchmark's "image channel"
attack surface.
"""

import logging
from pathlib import Path
from typing import Dict, Any, Optional, Generator, Iterable, List

try:
    from datasets import load_dataset
    from datasets import get_dataset_config_names
    HAS_DATASETS = True
except ImportError:
    HAS_DATASETS = False

from .base_adapter import BaseAdapter, DatasetMetadata

logger = logging.getLogger(__name__)


class MMSafetyBenchAdapter(BaseAdapter):
    """Adapter for PKU-Alignment/MM-SafetyBench (Extended; CC-BY-NC)."""

    DEFAULT_SPLITS = ["SD", "SD_TYPO", "TYPO"]

    def __init__(
        self,
        output_dir: Path,
        configs: Optional[List[str]] = None,
        splits: Optional[List[str]] = None,
        include_text_only: bool = False,
        export_images: bool = False,
        images_subdir: str = "mmsafetybench_images",
        verbose: bool = False,
    ):
        super().__init__(output_dir, verbose)

        if not HAS_DATASETS:
            raise ImportError("The 'datasets' library is required. Install with: pip install datasets")

        self.configs = configs
        self.splits = splits or list(self.DEFAULT_SPLITS)
        if include_text_only and "Text_only" not in self.splits:
            self.splits.append("Text_only")

        self.export_images = bool(export_images)
        self.images_dir_out = self.output_dir / images_subdir

        logger.warning(
            "MM-SafetyBench is CC-BY-NC licensed. Do not redistribute converted data or exported images."
        )

    @property
    def metadata(self) -> DatasetMetadata:
        return DatasetMetadata(
            name="MM-SafetyBench",
            version="1.0",
            license="CC-BY-NC-4.0",
            license_url="https://huggingface.co/datasets/PKU-Alignment/MM-SafetyBench",
            source_url="https://huggingface.co/datasets/PKU-Alignment/MM-SafetyBench",
            paper_url=None,
            can_redistribute=False,
            target_surface="B2",
            description="Multimodal safety benchmark prompts with images (research only)",
        )

    def _iter_configs(self) -> Iterable[str]:
        if self.configs:
            return self.configs
        try:
            return get_dataset_config_names("PKU-Alignment/MM-SafetyBench")
        except Exception as e:
            logger.warning(f"Failed to list MM-SafetyBench configs; using fallback list: {e}")
            return [
                "EconomicHarm",
                "Financial_Advice",
                "Fraud",
                "Gov_Decision",
                "HateSpeech",
                "Health_Consultation",
                "Illegal_Activitiy",
                "Legal_Opinion",
                "Malware_Generation",
                "Physical_Harm",
                "Political_Lobbying",
                "Privacy_Violence",
                "Sex",
            ]

    def load_source_data(self) -> Generator[Dict[str, Any], None, None]:
        if self.export_images:
            self.images_dir_out.mkdir(parents=True, exist_ok=True)

        for config in self._iter_configs():
            for split in self.splits:
                logger.info(f"Loading MM-SafetyBench config={config} split={split} ...")
                try:
                    dataset = load_dataset("PKU-Alignment/MM-SafetyBench", config, split=split)
                except Exception as e:
                    logger.error(f"Failed to load MM-SafetyBench config={config} split={split}: {e}")
                    continue

                for item in dataset:
                    rec = dict(item)
                    rec["_config"] = config
                    rec["_split"] = split
                    yield rec

    def transform_record(self, record: Dict[str, Any]) -> Optional[Dict[str, Any]]:
        question = record.get("question") or ""
        if not question or len(str(question).strip()) < 3:
            return None

        config = record.get("_config", "Unknown")
        split = record.get("_split", "Unknown")
        src_id = record.get("id")

        user_query = self.sanitize_text(str(question), max_length=10000)
        attack_id = self.generate_id("B2-MMSafetyBench", f"{config}:{split}:{src_id}:{user_query}")

        image_obj = record.get("image")
        exported_relpath = None
        if self.export_images and image_obj is not None:
            try:
                out_dir = self.images_dir_out / config / split
                out_dir.mkdir(parents=True, exist_ok=True)
                out_path = out_dir / f"{attack_id}.png"
                image_obj.save(out_path)
                exported_relpath = str(out_path.relative_to(self.output_dir))
            except Exception as e:
                logger.warning(f"Failed to export MM-SafetyBench image for {attack_id}: {e}")

        judge_pending = {
            "assessment": "Pending",
            "justification": "Imported from MM-SafetyBench; requires evaluation.",
            "confidence": "N/A",
        }

        return {
            "image_type": split.lower(),
            "attack_id": attack_id,
            "user_query": user_query,
            "adversarial_goal": user_query,
            "attack_subtype": f"MMSafetyBench_{config}_{split}",
            "poison_payload": "[IMAGE_TEXT]" if split in {"TYPO", "SD_TYPO"} else "",
            "baseline_output": "",
            "baseline_judge_assessment": judge_pending,
            "baseline_rag_success": True,
            "baseline_failure_stage": "Success",
            "mta_output": "",
            "mta_judge_assessment": judge_pending,
            "mta_rag_success": True,
            "mta_failure_stage": "Success",
            "script_version": self.BENCHMARK_VERSION,
            "source_dataset": "MM-SafetyBench",
            "source_id": str(src_id) if src_id is not None else None,
            "source_license": "CC-BY-NC-4.0 (Research Only)",
            "research_only": True,
            "hf_dataset": "PKU-Alignment/MM-SafetyBench",
            "hf_config": config,
            "hf_split": split,
            "mmsafetybench_image_output_relpath": exported_relpath,
        }


def main():
    import argparse

    parser = argparse.ArgumentParser(description="MM-SafetyBench to B2 Schema Adapter (Research Only)")
    parser.add_argument(
        "--configs",
        type=str,
        default=None,
        help="Comma-separated MM-SafetyBench configs to include (default: common set)",
    )
    parser.add_argument(
        "--splits",
        type=str,
        default=None,
        help="Comma-separated splits to include (default: SD,SD_TYPO,TYPO)",
    )
    parser.add_argument(
        "--include-text-only",
        action="store_true",
        help="Include the Text_only split (no image).",
    )
    parser.add_argument(
        "--export-images",
        action="store_true",
        help="Export images into output directory (CC-BY-NC; do not redistribute).",
    )
    parser.add_argument(
        "--output-dir",
        type=Path,
        default=Path("data/external_augmented/extended"),
        help="Output directory",
    )
    parser.add_argument(
        "--max-records",
        type=int,
        default=None,
        help="Maximum records to process",
    )
    parser.add_argument(
        "-v",
        "--verbose",
        action="store_true",
        help="Enable verbose output",
    )

    args = parser.parse_args()

    configs = [c.strip() for c in args.configs.split(",")] if args.configs else None
    splits = [s.strip() for s in args.splits.split(",")] if args.splits else None

    print("\n" + "=" * 60)
    print("WARNING: MM-SafetyBench is CC-BY-NC licensed.")
    print("Do not redistribute the converted data.")
    print("=" * 60 + "\n")

    adapter = MMSafetyBenchAdapter(
        output_dir=args.output_dir,
        configs=configs,
        splits=splits,
        include_text_only=args.include_text_only,
        export_images=args.export_images,
        verbose=args.verbose,
    )
    result = adapter.run(max_records=args.max_records)

    print("\nAdapter Results:")
    print(f"  Success: {result.success}")
    print(f"  Records processed: {result.records_processed}")
    print(f"  Records failed: {result.records_failed}")
    if result.output_file:
        print(f"  Output file: {result.output_file}")
    print("\nREMINDER: CC-BY-NC - non-commercial use only.")


if __name__ == "__main__":
    main()