eval-unlearn 1.0.0__py3-none-any.whl

eval_unlearn/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.1.0"

eval_unlearn/artifacts/__init__.py

@@ -0,0 +1,3 @@
+from .writer import ArtifactWriter
+
+__all__ = ["ArtifactWriter"]

eval_unlearn/artifacts/writer.py

@@ -0,0 +1,175 @@
+import os
+import json
+from typing import List, Any, Dict, Optional
+from ..logging_utils import get_logger
+
+logger = get_logger(__name__)
+
+
+class ArtifactWriter:
+    """
+    Handles saving of benchmark artifacts (images and reports).
+
+    Folder layout::
+
+        <base_dir>/
+          <technique>_<metric>_<run_id>/
+            images/
+              <category>/        ← only if metadata has categories
+                0.png, 1.png
+              0.png, 1.png       ← flat if no categories
+            <run_id>_report.json
+    """
+
+    def __init__(self, base_dir: str = "results"):
+        self.base_dir = base_dir
+
+    def save_run(
+        self,
+        run_id: str,
+        technique_name: str,
+        metric_name: str,
+        images: List[Any],
+        report: Optional[Dict[str, Any]] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+        detailed_report: Optional[Dict[str, Any]] = None,
+        image_index_offset: int = 0,
+        category_counters_init: Optional[Dict[str, int]] = None,
+    ) -> str:
+        """
+        Save images and report for a benchmark run.
+
+        Args:
+            run_id: Short hash identifying this run.
+            technique_name: Name of the technique (e.g. ``"sld"``).
+            metric_name: Name of the metric (e.g. ``"asr"``).
+            images: List of generated PIL images.
+            report: Simplified result dictionary to persist as ``{run_id}_report.json``.
+                If None, no report is saved.
+            metadata: Dataset metadata. If it contains a ``categories``
+                key (list parallel to *images*), images are saved into
+                per-category subdirectories.
+            detailed_report: Extended result dictionary including technique and metric
+                configs, saved alongside the simplified report as
+                ``{run_id}_report_full.json``. If None, no detailed report is saved.
+            image_index_offset: Starting index for flat image filenames. Pass
+                ``total_generated`` before the current batch to avoid overwriting
+                images saved in earlier batches.
+            category_counters_init: Per-category filename counters accumulated from
+                previous batches, used to continue numbering in category-aware saving.
+
+        Returns:
+            Path to the saved simplified report JSON (or where it would be saved).
+        """
+        metadata = metadata or {}
+        categories = metadata.get("categories")
+
+        image_paths = []
+
+        # If images provided, save them in structured folder
+        if images:
+            # Build folder: <base_dir>/<technique>_<metric>_<run_id>/
+            folder_name = f"{technique_name}_{metric_name}_{run_id}"
+            run_dir = os.path.join(self.base_dir, folder_name)
+            images_dir = os.path.join(run_dir, "images")
+            os.makedirs(images_dir, exist_ok=True)
+
+            # Save images — category-aware if metadata provides categories
+            logger.info(f"Saving {len(images)} images to {images_dir}...")
+
+            if categories and len(categories) == len(images):
+                # Category-aware: save into subdirectories
+                category_counters: Dict[str, int] = dict(category_counters_init or {})
+                for img, cat in zip(images, categories):
+                    cat_dir = os.path.join(images_dir, cat.lower())
+                    os.makedirs(cat_dir, exist_ok=True)
+                    idx = category_counters.get(cat.lower(), 0)
+                    category_counters[cat.lower()] = idx + 1
+                    path = os.path.join(cat_dir, f"{technique_name}_{metric_name}_{run_id}_{idx}.png")
+                    image_paths.append(self._save_image(img, path, idx))
+            else:
+                # Flat: save numbered images directly
+                for i, img in enumerate(images):
+                    global_i = i + image_index_offset
+                    path = os.path.join(images_dir, f"{technique_name}_{metric_name}_{run_id}_{global_i}.png")
+                    image_paths.append(self._save_image(img, path, global_i))
+
+            # Filter out None (failed saves)
+            image_paths = [p for p in image_paths if p is not None]
+        else:
+            logger.info("No images to save")
+
+        # Resolve the directory where reports are saved
+        if images:
+            folder_name = f"{technique_name}_{metric_name}_{run_id}"
+            report_dir = os.path.join(self.base_dir, folder_name)
+        else:
+            os.makedirs(self.base_dir, exist_ok=True)
+            report_dir = self.base_dir
+
+        report_path = os.path.join(report_dir, f"{run_id}_report.json")
+
+        # Save simplified report
+        if report is not None:
+            try:
+                with open(report_path, "w") as f:
+                    json.dump(report, f, indent=4)
+                logger.info(f"Report saved to {report_path}")
+            except Exception as e:
+                logger.error(f"Failed to save report: {e}")
+        else:
+            logger.info("Skipping report save (not provided)")
+
+        # Save detailed report
+        if detailed_report is not None:
+            detailed_path = os.path.join(report_dir, f"{run_id}_report_full.json")
+            try:
+                with open(detailed_path, "w") as f:
+                    json.dump(detailed_report, f, indent=4)
+                logger.info(f"Detailed report saved to {detailed_path}")
+            except Exception as e:
+                logger.error(f"Failed to save detailed report: {e}")
+
+        if report is not None:
+            self._sync_to_final_reports(report_path, report)
+
+        return report_path
+
+    def _sync_to_final_reports(self, report_path: str, report: Dict[str, Any]) -> None:
+        """Copy report as {technique}_{concept}_latest_report.json in base_dir if newer."""
+        technique = report.get("technique_name", "unknown")
+        concept = report.get("erase_concept", "unknown")
+        timestamp = report.get("timestamp", 0)
+
+        dest = os.path.join(self.base_dir, f"{technique}_{concept}_latest_report.json")
+
+        if os.path.exists(dest):
+            try:
+                with open(dest) as f:
+                    existing = json.load(f)
+                if existing.get("timestamp", 0) >= timestamp:
+                    logger.info(f"Skipping latest report sync — existing is newer or equal: {dest}")
+                    return
+            except Exception:
+                pass
+
+        try:
+            import shutil
+            shutil.copy2(report_path, dest)
+            logger.info(f"Synced latest report to {dest}")
+        except Exception as e:
+            logger.error(f"Failed to sync latest report: {e}")
+
+    @staticmethod
+    def _save_image(img: Any, path: str, index: int) -> Optional[str]:
+        """Save a single image, returning the path or None on failure."""
+        try:
+            if hasattr(img, "save"):
+                img.save(path)
+                return path
+            else:
+                logger.warning(f"Image {index} does not support save(), skipping.")
+                return None
+        except Exception as e:
+            logger.error(f"Failed to save image {index}: {e}")
+            return None

eval_unlearn/cli.py

@@ -0,0 +1,329 @@
+import argparse
+import sys
+import os
+import json
+import logging
+import warnings
+from typing import Any, Dict, List
+
+from dotenv import load_dotenv
+
+load_dotenv(override=True)
+
+# Suppress noisy library output before any heavy imports
+logging.getLogger("transformers").setLevel(logging.ERROR)
+logging.getLogger("datasets").setLevel(logging.ERROR)
+logging.getLogger("huggingface_hub").setLevel(logging.ERROR)
+logging.getLogger("onnxruntime").setLevel(logging.ERROR)
+logging.getLogger("tensorflow").setLevel(logging.ERROR)
+warnings.filterwarnings("ignore")
+
+from eval_unlearn.runners import SingleBenchmarkRunner, MultiBenchmarkRunner
+from eval_unlearn.logging_utils import get_logger
+
+logger = get_logger("cli")
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+def _load_config(path: str) -> Dict[str, Any]:
+    if not os.path.exists(path):
+        logger.error("Config file not found: %s", path)
+        sys.exit(1)
+    with open(path) as f:
+        if path.endswith((".yaml", ".yml")):
+            try:
+                import yaml
+                return yaml.safe_load(f)
+            except ImportError:
+                logger.error("PyYAML not installed. Use a .json config or: pip install pyyaml")
+                sys.exit(1)
+        return json.load(f)
+
+
+def _parse_metrics_list(metrics_list: List[Dict[str, Any]]):
+    metric_names, metric_configs = [], {}
+    for m in metrics_list:
+        name = m.get("name")
+        if not name:
+            logger.error("Each entry in 'metrics' must have a 'name'")
+            sys.exit(1)
+        metric_names.append(name)
+        cfg = m.get("config", {})
+        if cfg:
+            metric_configs[name] = cfg
+    return metric_names, metric_configs
+
+
+def _build_single_runner(config: Dict[str, Any], output_dir: str) -> SingleBenchmarkRunner:
+    tech = config.get("technique", {})
+    tech_name = tech.get("name")
+    if not tech_name:
+        logger.error("Config must specify 'technique.name'")
+        sys.exit(1)
+    met = config.get("metric", {})
+    met_name = met.get("name")
+    if not met_name:
+        logger.error("Config must specify 'metric.name'")
+        sys.exit(1)
+    logger.info("Mode: single | technique=%s | metric=%s", tech_name, met_name)
+    return SingleBenchmarkRunner(
+        technique_name=tech_name,
+        metric_name=met_name,
+        technique_config=tech.get("config", {}),
+        metric_config=met.get("config", {}),
+        output_dir=output_dir,
+    )
+
+
+def _build_multi_runner(config: Dict[str, Any], output_dir: str) -> MultiBenchmarkRunner:
+    tech = config.get("technique", {})
+    tech_name = tech.get("name")
+    if not tech_name:
+        logger.error("Config must specify 'technique.name'")
+        sys.exit(1)
+    metrics_list = config.get("metrics", [])
+    if not metrics_list:
+        logger.error("Config must specify 'metrics' as a non-empty list")
+        sys.exit(1)
+    metric_names, metric_configs = _parse_metrics_list(metrics_list)
+    logger.info("Mode: multi | technique=%s | metrics=%s", tech_name, metric_names)
+    return MultiBenchmarkRunner(
+        technique_name=tech_name,
+        metric_names=metric_names,
+        technique_config=tech.get("config", {}),
+        metric_configs=metric_configs,
+        output_dir=output_dir,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Commands
+# ---------------------------------------------------------------------------
+
+def cmd_run(args):
+    """Execute a benchmark run, then optionally push results to HF Hub."""
+    config = _load_config(args.config)
+    output_dir = config.get("output_dir", "results")
+
+    has_technique = "technique" in config
+    has_metrics = "metrics" in config
+    has_metric = "metric" in config
+
+    try:
+        if has_technique and has_metrics:
+            runner = _build_multi_runner(config, output_dir)
+        elif has_technique and has_metric:
+            runner = _build_single_runner(config, output_dir)
+        else:
+            logger.error(
+                "Invalid config: must have 'technique'+'metric' (single run) "
+                "or 'technique'+'metrics' (multi-metric run)"
+            )
+            sys.exit(1)
+    except ValueError as e:
+        logger.error(str(e))
+        sys.exit(1)
+
+    try:
+        report = runner.run()
+    except Exception:
+        logger.exception("Run failed.")
+        sys.exit(1)
+
+    logger.info("Run completed. run_id=%s", report.get("run_id"))
+
+    if args.hf_repo:
+        _push_results(
+            output_dir=output_dir,
+            hf_repo=args.hf_repo,
+            hf_path=args.hf_path,
+            create_pr=args.create_pr,
+        )
+
+
+def _push_results(output_dir: str, hf_repo: str, hf_path: str, create_pr: bool):
+    from eval_unlearn.hub import HFSync
+    remote_path = hf_path or os.path.basename(os.path.normpath(output_dir))
+    sync = HFSync(repo_id=hf_repo, create_pr=create_pr)
+    try:
+        url = sync.push_folder(output_dir, remote_path)
+        logger.info("Results pushed → %s/%s  (%s)", hf_repo, remote_path, url)
+    except Exception as e:
+        logger.error("HF push failed: %s", e)
+        sys.exit(1)
+
+
+def cmd_push(args):
+    """Push a local directory to HF Hub."""
+    from eval_unlearn.hub import HFSync
+    remote_path = args.remote_path or os.path.basename(os.path.normpath(args.local_dir))
+    sync = HFSync(repo_id=args.repo, create_pr=args.create_pr)
+    try:
+        url = sync.push_folder(args.local_dir, remote_path)
+        logger.info("Pushed %s → %s/%s  (%s)", args.local_dir, args.repo, remote_path, url)
+    except Exception as e:
+        logger.error("Push failed: %s", e)
+        sys.exit(1)
+
+
+def cmd_pull(args):
+    """Pull artifacts from HF Hub."""
+    from eval_unlearn.hub import HFSync
+    local_dir = args.local_dir or "results"
+    sync = HFSync(repo_id=args.repo)
+    try:
+        if args.remote_path:
+            path = sync.pull_folder(args.remote_path, local_dir)
+        else:
+            path = sync.pull_all(local_dir)
+        logger.info("Downloaded to %s", path)
+    except Exception as e:
+        logger.error("Pull failed: %s", e)
+        sys.exit(1)
+
+
+def cmd_plugins(_args):
+    """List all registered techniques, metrics, and datasets."""
+    from eval_unlearn.registry.entrypoints import load_entrypoints
+    from eval_unlearn.registry.local import _TECHNIQUES, _METRICS, _DATASETS
+
+    load_entrypoints()
+
+    print("\nTechniques:")
+    for name in sorted(_TECHNIQUES):
+        print(f"  {name}")
+    print("\nMetrics:")
+    for name in sorted(_METRICS):
+        print(f"  {name}")
+    print("\nDatasets:")
+    for name in sorted(_DATASETS):
+        print(f"  {name}")
+    print()
+
+
+def cmd_models(_args):
+    """Show the base model used by each technique and metric."""
+    from eval_unlearn.techniques._base_models import TECHNIQUE_BASE_MODELS
+    from eval_unlearn.metrics._base_models import METRIC_MODELS
+
+    # Techniques
+    print("\nTechniques:")
+    print(f"  {'name':<20} {'base model':<45} {'configurable'}")
+    print(f"  {'-'*20} {'-'*45} {'-'*12}")
+    for name in sorted(TECHNIQUE_BASE_MODELS):
+        print(f"  {name:<20} {TECHNIQUE_BASE_MODELS[name]:<45} no")
+    print(f"  {'free_run':<20} {'(user-specified via model_id)':<45} required")
+
+    # Metrics
+    print("\nMetrics:")
+    print(f"  {'name':<20} {'model':<45} {'configurable'}")
+    print(f"  {'-'*20} {'-'*45} {'-'*12}")
+    for name in sorted(METRIC_MODELS):
+        info = METRIC_MODELS[name]
+        if info.configurable:
+            if info.choices:
+                choices_str = " | ".join(c.split("/")[-1] for c in sorted(info.choices))
+                detail = f"yes  (config: {info.config_field}: {choices_str})"
+            else:
+                detail = f"yes  (config: {info.config_field})"
+        elif info.note:
+            detail = f"no   ({info.note})"
+        else:
+            detail = "no"
+        print(f"  {name:<20} {info.model:<45} {detail}")
+    print()
+
+
+# ---------------------------------------------------------------------------
+# Argument parser
+# ---------------------------------------------------------------------------
+
+def main():
+    parser = argparse.ArgumentParser(
+        prog="eval-unlearn",
+        description="Eval-Learn: unlearning benchmark for text-to-image diffusion models",
+    )
+    parser.add_argument("--version", action="store_true", help="Print version and exit")
+    subparsers = parser.add_subparsers(dest="command")
+
+    # run
+    run_p = subparsers.add_parser("run", help="Execute a benchmark run")
+    run_p.add_argument("--config", "-c", required=True, help="Path to config file (JSON/YAML)")
+    run_p.add_argument(
+        "--hf-repo",
+        default=None,
+        metavar="REPO_ID",
+        help="HF Hub dataset repo to push results to after the run (e.g. org/my-results)",
+    )
+    run_p.add_argument(
+        "--hf-path",
+        default=None,
+        metavar="REMOTE_PATH",
+        help="Remote path inside the repo (default: basename of output_dir from config)",
+    )
+    run_p.add_argument(
+        "--create-pr",
+        action="store_true",
+        help="Open a pull request instead of committing directly to HF Hub",
+    )
+
+    # push
+    push_p = subparsers.add_parser("push", help="Push a local results directory to HF Hub")
+    push_p.add_argument("--repo", required=True, metavar="REPO_ID", help="HF Hub dataset repo ID")
+    push_p.add_argument("--local-dir", required=True, metavar="PATH", help="Local directory to upload")
+    push_p.add_argument(
+        "--remote-path",
+        default=None,
+        metavar="PATH",
+        help="Destination path in the repo (default: basename of --local-dir)",
+    )
+    push_p.add_argument("--create-pr", action="store_true")
+
+    # pull
+    pull_p = subparsers.add_parser("pull", help="Pull artifacts from HF Hub")
+    pull_p.add_argument("--repo", required=True, metavar="REPO_ID", help="HF Hub dataset repo ID")
+    pull_p.add_argument(
+        "--remote-path",
+        default=None,
+        metavar="PATH",
+        help="Remote path to download (omit to pull the entire repo)",
+    )
+    pull_p.add_argument(
+        "--local-dir",
+        default=None,
+        metavar="PATH",
+        help="Local directory to download into (default: results/)",
+    )
+
+    # plugins
+    subparsers.add_parser("plugins", help="List all registered techniques, metrics, and datasets")
+
+    # models
+    subparsers.add_parser("models", help="Show the base model used by each technique and metric")
+
+    args = parser.parse_args()
+
+    if args.version:
+        from eval_unlearn import __version__
+        print(f"eval-unlearn {__version__}")
+        sys.exit(0)
+
+    dispatch = {
+        "run": cmd_run,
+        "push": cmd_push,
+        "pull": cmd_pull,
+        "plugins": cmd_plugins,
+        "models": cmd_models,
+    }
+
+    if args.command in dispatch:
+        dispatch[args.command](args)
+    else:
+        parser.print_help()
+
+
+if __name__ == "__main__":
+    main()

eval_unlearn/config/hf_datasets.yaml

@@ -0,0 +1,37 @@
+# HuggingFace dataset paths and column configs.
+# Loads entire datasets without split specification. HF_TOKEN should be set in your .env file.
+
+datasets:
+
+  coco:
+    repo_id: "phiyodr/coco2017"
+    split: "train"
+    caption_col: "captions"
+    url_col: "coco_url"
+
+  i2p:
+    repo_id: "AIML-TUDA/i2p"
+    split: "train"
+    caption_col: "prompt"
+    concept_col: "categories"
+
+  tifa:
+    repo_id: "Unlearningltd/datasets"
+    data_files: "tifa/tifa_dataset.csv"
+    split: "train"
+    caption_col: "caption"
+    qa_col: "qas"
+
+  err_challenge:
+    repo_id: "Unlearningltd/datasets"
+    data_files: "ERR/raw_csv_data/challenge_dataset.csv"
+    split: "train"
+    caption_col: "direct_prompt"
+    concept_col: "concept_name"
+
+  ring_a_bell:
+    repo_id: "Unlearningltd/datasets"
+    data_files: "ring_a_bell/ring_a_bell_dataset.csv"
+    split: "train"
+    caption_col: "prompt"
+    concept_col: "concept"

eval_unlearn/configs/__init__.py

@@ -0,0 +1,3 @@
+from .base import BaseConfig
+
+__all__ = ["BaseConfig"]

eval_unlearn/configs/base.py

@@ -0,0 +1,23 @@
+from dataclasses import dataclass, asdict
+from typing import Dict, Any
+
+
+@dataclass(frozen=True)
+class BaseConfig:
+    """
+    Base configuration class providing dictionary serialization methods.
+    """
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert the configuration to a dictionary."""
+        return asdict(self)
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "BaseConfig":
+        """
+        Create a configuration instance from a dictionary.
+        ignores keys that do not match fields in the dataclass to ensure forward compatibility.
+        """
+        valid_keys = {k for k, f in cls.__dataclass_fields__.items() if f.init}
+        filtered_data = {k: v for k, v in data.items() if k in valid_keys}
+        return cls(**filtered_data)

eval_unlearn/datasets/__init__.py

@@ -0,0 +1,13 @@
+from .coco_parquet import load_coco_parquet
+from .i2p_csv import load_i2p_csv
+from .err_composite import load_err_composite
+from .tifa_csv import load_tifa_csv
+from .ua_ira_csv import load_ua_ira_csv
+
+__all__ = [
+    "load_coco_parquet",
+    "load_i2p_csv",
+    "load_err_composite",
+    "load_tifa_csv",
+    "load_ua_ira_csv",
+]