Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions src/anonymizer/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
__version__ = version("nemo-anonymizer")

from data_designer.config.models import ModelProvider as ModelProvider
from data_designer.config.run_config import RunConfig as RunConfig

from anonymizer.config.anonymizer_config import (
AnonymizerConfig,
Expand Down Expand Up @@ -64,6 +65,7 @@ def __getattr__(name: str) -> object:
"Redact",
"Rewrite",
"RiskTolerance",
"RunConfig",
"Substitute",
"configure_logging",
]
8 changes: 8 additions & 0 deletions src/anonymizer/interface/anonymizer.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@
from typing import TYPE_CHECKING

from data_designer.config.models import ModelProvider
from data_designer.config.run_config import RunConfig
from data_designer.config.utils.io_helpers import load_config_file
from data_designer.interface.data_designer import DataDesigner

Expand Down Expand Up @@ -109,6 +110,7 @@ def __init__(
model_providers: list[ModelProvider] | str | Path | None = None,
artifact_path: str | Path | None = None,
data_designer: DataDesigner | None = None,
data_designer_run_config: RunConfig | None = None,
detection_workflow: EntityDetectionWorkflow | None = None,
replace_runner: ReplacementWorkflow | None = None,
rewrite_runner: RewriteWorkflow | None = None,
Expand All @@ -125,6 +127,10 @@ def __init__(
artifact_path: Directory for intermediate artifacts. Defaults to
``.anonymizer-artifacts``.
data_designer: Pre-configured DataDesigner instance (advanced usage).
data_designer_run_config: Optional DataDesigner run configuration
applied to the Anonymizer-managed or caller-supplied DataDesigner
instance. Use this for DataDesigner execution knobs such as
``buffer_size`` and ``max_in_flight_tasks``.
detection_workflow: Custom detection workflow (advanced/testing).
replace_runner: Custom replacement workflow (advanced/testing).
rewrite_runner: Custom rewrite workflow (advanced/testing).
Expand Down Expand Up @@ -164,6 +170,8 @@ def __init__(
model_providers=self._resolved_providers,
)
reapply_log_levels()
if data_designer_run_config is not None:
self._data_designer.set_run_config(data_designer_run_config)
self._adapter = NddAdapter(data_designer=self._data_designer)
self._detection_workflow = detection_workflow or EntityDetectionWorkflow(adapter=self._adapter)
self._replace_runner = replace_runner or ReplacementWorkflow(
Expand Down
32 changes: 32 additions & 0 deletions tests/interface/test_anonymizer_interface.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
import pytest
from data_designer.config.models import ModelConfig

from anonymizer import RunConfig
from anonymizer.config.anonymizer_config import AnonymizerConfig, AnonymizerInput, Rewrite
from anonymizer.config.models import ModelSelection, ReplaceModelSelection
from anonymizer.config.replace_strategies import Redact, Substitute
Expand Down Expand Up @@ -169,6 +170,37 @@ def test_anonymizer_custom_model_providers_override_bundled_defaults() -> None:
assert passed_providers is custom_providers


def test_anonymizer_applies_data_designer_run_config_to_managed_instance() -> None:
run_config = RunConfig(buffer_size=20, max_in_flight_tasks=64)

with patch("anonymizer.interface.anonymizer.DataDesigner") as mock_data_designer:
Anonymizer(
data_designer_run_config=run_config,
detection_workflow=Mock(),
replace_runner=Mock(),
rewrite_runner=Mock(),
)

mock_data_designer.return_value.set_run_config.assert_called_once_with(run_config)


def test_anonymizer_applies_data_designer_run_config_to_supplied_instance() -> None:
from data_designer.interface.data_designer import DataDesigner

data_designer = Mock(spec=DataDesigner)
run_config = RunConfig(buffer_size=20, max_in_flight_tasks=64)

Anonymizer(
data_designer=data_designer,
data_designer_run_config=run_config,
detection_workflow=Mock(),
replace_runner=Mock(),
rewrite_runner=Mock(),
)

data_designer.set_run_config.assert_called_once_with(run_config)


def test_anonymizer_rejects_missing_provider_as_invalid_config_error() -> None:
yaml_str = """
model_configs:
Expand Down
Loading