diff --git a/plans/237/combined-rewrite-graph.md b/plans/237/combined-rewrite-graph.md new file mode 100644 index 00000000..07a09369 --- /dev/null +++ b/plans/237/combined-rewrite-graph.md @@ -0,0 +1,123 @@ + + + +# Combined Rewrite Integration + +Tracks the proof of concept for [GitHub issue #237](https://github.com/NVIDIA-NeMo/Anonymizer/issues/237). +It extends the rewrite portion of [Anonymizer Workflow Columns](../custom-column-plugins/anonymizer-workflow-columns.md). + +## Goal + +Replace the Python-controlled post-detection rewrite loop with one DataDesigner +execution: + +```text +replacement map + -> domain + disposition + QA + initial rewrite + -> evaluate 0 + -> repair 0 when evaluate 0 fails + -> evaluate 1 when repair 0 ran + -> ... + -> coalesce the last executed evaluation state +``` + +The repair count is static at graph-build time. Each configured round gets unique +columns and uses `SkipConfig` to bypass repair and downstream re-evaluation for +rows that already pass. This PR keeps the legacy workflow as the default. After +[Data Designer #861](https://github.com/NVIDIA-NeMo/DataDesigner/pull/861) is +merged and released, a small follow-up can consume its terminal-failure API, +make the combined graph the default, and retain legacy as an explicit fallback. + +## Current Status + +`CombinedRewriteWorkflow` remains an opt-in `RewriteWorkflow` subclass while the +legacy path serves as the default and parity oracle. This avoids losing precise +failure-stage attribution before Data Designer #861 is available in a supported +release and Anonymizer maps its failed column and seed-row position back to the +existing `FailedRecord` contract. + +The graph currently: + +- generates and filters the replacement map in the same execution; +- reuses the existing domain, disposition, QA, rewrite, evaluator, and repair helpers; +- supports `max_repair_iterations >= 0` by statically unrolling rounds; +- preserves no-entity passthrough outside DataDesigner; +- restores the existing final rewrite, metric, repair-count, and review columns; +- leaves the separate `evaluate()` judge path unchanged. + +With two repair rounds, the graph has 36 columns and DataDesigner 0.8 validates it +without duplicate producers, missing dependencies, or cycles. + +Tests execute real Data Designer conditional scheduling for mixed rows requiring +zero, one, two, and more-than-allowed repairs. They also cover no-entity +passthrough, mixed-row ordering, final state selection, repair counts, exhausted +review flags, malformed initial rewrites, coarse combined-boundary failures, and +graph validation with up to ten repair rounds. Local 64-row batches cover both +mostly-skipped and mostly-repaired scheduling. + +## Expected Execution Change + +For a rewrite run with entity rows: + +| Path | Base DD runs | DD runs per repair round | +|---|---:|---:| +| Current full pipeline | 5 | 2 | +| Proof of concept | 3 | 0 | + +The totals include the two existing detection runs. This proof combines only the +post-detection rewrite work, so a full detection-plus-rewrite graph remains a later step. + +## Benchmark Result + +The authoritative controlled run completed 12 pairs without workload failures. +Both paths received the same prepared and initially evaluated state. Ten rows +required one repair and two skipped repair in both paths; none exhausted the +three configured rounds. Both paths made 80 LLM requests, while the measured +rewrite/evaluate portion used four Data Designer workflows for legacy and two +for combined across two counterbalanced groups. + +Combined took 10.56 seconds versus 10.80 seconds for legacy, a 2.2% reduction. +Leakage remained zero for every output, repair counts agreed for all rows, 10/12 +outputs were byte-identical, and 11/12 review decisions agreed. The paired mean +utility delta was -0.0867 with median zero and an approximate 95% interval of +[-0.2496, 0.0763]. Separate real-model generation and judge calls remain +nondeterministic, so this establishes latency and behavioral parity rather than +a speedup. + +Performance is therefore a regression guardrail, not the integration rationale. + +## Completion and Follow-up + +1. [x] **Conditional behavior**: verify row-local skipping, multiple repairs, + exhausted repairs, passthrough defaults, row order, and graph validation. +2. [x] **Failure-attribution scope**: keep this implementation opt-in until the + terminal-failure API in + [Data Designer #861](https://github.com/NVIDIA-NeMo/DataDesigner/pull/861) + is merged, released, and integrated. Data Designer 0.8 task traces expose + column and row position only with full tracing, which is not an appropriate + production result contract. +3. [x] **Measurements**: record one physical `rewrite-combined` Data Designer + workflow while preserving aggregate model usage, repair counts, review flags, + and runner-level row counts. Precise failure-stage measurements remain part of + the failure-attribution gate. +4. [x] **Behavioral equivalence**: compare legacy and combined public outputs with + deterministic repaired results, and cover partial row loss and malformed + initial rewrites. +5. [x] **Scale guardrails**: local 64-row mostly-skipped and mostly-repaired + batches pass, an earlier Slurm suite completed 60 executions without workload + failures, and retained GB300 telemetry showed flat HBM use. Combined artifacts + were larger but remained below 230 KB in the controlled runs. The retained + telemetry cadence does not support tail-latency conclusions, which are not + required for this opt-in proof of concept. +6. [x] **Default rollout deferred**: a follow-up after Data Designer #861 will + consume terminal failure provenance, make the combined graph the default, and + retain the legacy path as an explicit fallback. Legacy removal can follow + production rollout evidence. +7. [x] **Performance guardrail**: rerun the corrected paired benchmark with + balanced ordering and equivalent repair decisions. + +## Portability + +The current closure-based custom columns are sufficient for the local in-process +path. Serializable plugin configs remain part of the broader workflow-column plan +and become a prerequisite when distributed rewrite graph export is supported. diff --git a/src/anonymizer/config/anonymizer_config.py b/src/anonymizer/config/anonymizer_config.py index 5a7c8d66..3afdea1c 100644 --- a/src/anonymizer/config/anonymizer_config.py +++ b/src/anonymizer/config/anonymizer_config.py @@ -131,6 +131,10 @@ class Rewrite(BaseModel): ge=0, description="Maximum repair rounds. Set to 0 to disable repair.", ) + use_combined_graph: bool = Field( + default=False, + description="Run rewrite and conditional repair iterations in one Data Designer graph.", + ) strict_entity_protection: bool = Field( default=False, description="If True, requires every entity to receive a protective disposition during sensitivity analysis.", diff --git a/src/anonymizer/engine/constants.py b/src/anonymizer/engine/constants.py index 18295caa..d9c51f95 100644 --- a/src/anonymizer/engine/constants.py +++ b/src/anonymizer/engine/constants.py @@ -45,6 +45,7 @@ COL_ENTITIES_BY_VALUE = "_entities_by_value" COL_REPLACED_TEXT = "__nemo_anonymizer_text_output__" COL_REPLACEMENT_MAP = "_replacement_map" +COL_REPLACEMENT_MAP_RAW = COL_REPLACEMENT_MAP + "__raw" COL_REPLACEMENT_MAP_SOURCE = "_replacement_map_source" # LlmReplaceWorkflow internal prompt-construction columns. Created by @@ -109,6 +110,7 @@ COL_QUALITY_QA = "_quality_qa" COL_PRIVACY_QA = "_privacy_qa" COL_REWRITTEN_TEXT = "_rewritten_text" # pre-repair intermediate; renamed to {text_col}_rewritten in user output +COL_REWRITTEN_TEXT_INITIAL = COL_REWRITTEN_TEXT + "__initial" COL_QUALITY_QA_REANSWER = "_quality_qa_reanswer" COL_QUALITY_QA_COMPARE = "_quality_qa_compare" COL_PRIVACY_QA_REANSWER = "_privacy_qa_reanswer" diff --git a/src/anonymizer/engine/rewrite/combined_workflow.py b/src/anonymizer/engine/rewrite/combined_workflow.py new file mode 100644 index 00000000..73e39f11 --- /dev/null +++ b/src/anonymizer/engine/rewrite/combined_workflow.py @@ -0,0 +1,534 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +from dataclasses import dataclass +from typing import Any + +import pandas as pd +from data_designer.config import SkipConfig, custom_column_generator +from data_designer.config.column_configs import CustomColumnConfig, LLMStructuredColumnConfig +from data_designer.config.column_types import ColumnConfigT +from data_designer.config.models import ModelConfig +from pydantic import BaseModel + +from anonymizer.config.models import ReplaceModelSelection, RewriteModelSelection +from anonymizer.config.rewrite import EvaluationCriteria, PrivacyGoal +from anonymizer.engine.constants import ( + COL_ANY_HIGH_LEAKED, + COL_ENTITIES_BY_VALUE, + COL_ENTITIES_FOR_REPLACE, + COL_ENTITIES_FOR_REPLACE_JSON, + COL_ENTITY_EXAMPLES, + COL_LEAKAGE_MASS, + COL_LEAKED_PRIVACY_ITEMS, + COL_NEEDS_HUMAN_REVIEW, + COL_NEEDS_REPAIR, + COL_PRIVACY_QA_REANSWER, + COL_QUALITY_QA_COMPARE, + COL_QUALITY_QA_REANSWER, + COL_REPAIR_ITERATIONS, + COL_REPLACEMENT_MAP, + COL_REPLACEMENT_MAP_RAW, + COL_REPLACEMENT_MAP_SOURCE, + COL_REWRITTEN_TEXT, + COL_REWRITTEN_TEXT_INITIAL, + COL_REWRITTEN_TEXT_NEXT, + COL_UTILITY_SCORE, + COL_WEIGHTED_LEAKAGE_RATE, +) +from anonymizer.engine.ndd.adapter import RECORD_ID_COLUMN, NddAdapter +from anonymizer.engine.ndd.model_loader import resolve_model_alias +from anonymizer.engine.replace.llm_replace_workflow import ( + REPLACEMENT_MAP_SOURCE_LLM, + _create_entity_examples, + _enrich_entities_for_template, + _filter_replacement_map_to_input_entities, + _get_replacement_mapping_prompt, +) +from anonymizer.engine.rewrite.domain_classification import DomainClassificationWorkflow +from anonymizer.engine.rewrite.evaluate import EvaluateWorkflow +from anonymizer.engine.rewrite.parsers import normalize_payload +from anonymizer.engine.rewrite.qa_generation import QAGenerationWorkflow +from anonymizer.engine.rewrite.repair import RepairWorkflow +from anonymizer.engine.rewrite.rewrite_generation import RewriteGenerationWorkflow +from anonymizer.engine.rewrite.rewrite_workflow import ( + RewriteResult, + RewriteWorkflow, + _apply_passthrough_defaults, + _has_entities, + _join_new_columns, +) +from anonymizer.engine.rewrite.sensitivity_disposition import SensitivityDispositionWorkflow +from anonymizer.engine.rewrite.workflow_utils import derive_seed_columns, select_seed_cols +from anonymizer.engine.row_partitioning import merge_and_reorder, split_rows +from anonymizer.engine.schemas import EntitiesByValueSchema, EntityReplacementMapSchema +from anonymizer.measurement import stage_timer + + +@dataclass(frozen=True) +class EvaluationState: + iteration: int + rewritten_text: str + quality_reanswer: str + privacy_reanswer: str + quality_compare: str + utility_score: str + leakage_mass: str + weighted_leakage_rate: str + any_high_leaked: str + needs_repair: str + + +@dataclass(frozen=True) +class RepairState: + iteration: int + leaked_items: str + rewritten_text: str + + +@dataclass(frozen=True) +class CombinedRewriteGraph: + columns: list[ColumnConfigT] + evaluation_states: list[EvaluationState] + repair_states: list[RepairState] + internal_columns: list[str] + + +class _FinalizationParams(BaseModel): + flag_utility_below: float | None + flag_leakage_above: float | None + + +def _iteration_column(column: str, iteration: int) -> str: + return f"{column}__iteration_{iteration}" + + +def _evaluation_state(iteration: int, rewritten_text: str) -> EvaluationState: + return EvaluationState( + iteration=iteration, + rewritten_text=rewritten_text, + quality_reanswer=_iteration_column(COL_QUALITY_QA_REANSWER, iteration), + privacy_reanswer=_iteration_column(COL_PRIVACY_QA_REANSWER, iteration), + quality_compare=_iteration_column(COL_QUALITY_QA_COMPARE, iteration), + utility_score=_iteration_column(COL_UTILITY_SCORE, iteration), + leakage_mass=_iteration_column(COL_LEAKAGE_MASS, iteration), + weighted_leakage_rate=_iteration_column(COL_WEIGHTED_LEAKAGE_RATE, iteration), + any_high_leaked=_iteration_column(COL_ANY_HIGH_LEAKED, iteration), + needs_repair=_iteration_column(COL_NEEDS_REPAIR, iteration), + ) + + +def _repair_state(iteration: int) -> RepairState: + return RepairState( + iteration=iteration, + leaked_items=_iteration_column(COL_LEAKED_PRIVACY_ITEMS, iteration), + rewritten_text=_iteration_column(COL_REWRITTEN_TEXT_NEXT, iteration), + ) + + +@custom_column_generator( + required_columns=[COL_ENTITIES_BY_VALUE], + side_effect_columns=[COL_ENTITIES_FOR_REPLACE, COL_ENTITIES_FOR_REPLACE_JSON], +) +def _prepare_replacement_inputs(row: dict[str, Any]) -> dict[str, Any]: + parsed = EntitiesByValueSchema.from_raw(row.get(COL_ENTITIES_BY_VALUE)) + row[COL_ENTITY_EXAMPLES] = _create_entity_examples(parsed) + row[COL_ENTITIES_FOR_REPLACE] = _enrich_entities_for_template(parsed) + row[COL_ENTITIES_FOR_REPLACE_JSON] = json.dumps(row[COL_ENTITIES_FOR_REPLACE]) + return row + + +@custom_column_generator( + required_columns=[COL_REPLACEMENT_MAP_RAW, COL_ENTITIES_BY_VALUE], + side_effect_columns=[COL_REPLACEMENT_MAP_SOURCE], +) +def _filter_replacement_map(row: dict[str, Any]) -> dict[str, Any]: + row[COL_REPLACEMENT_MAP] = _filter_replacement_map_to_input_entities( + raw_map=row.get(COL_REPLACEMENT_MAP_RAW, {"replacements": []}), + parsed_entities=EntitiesByValueSchema.from_raw(row.get(COL_ENTITIES_BY_VALUE)), + record_id=str(row.get(RECORD_ID_COLUMN, "")), + ) + row[COL_REPLACEMENT_MAP_SOURCE] = REPLACEMENT_MAP_SOURCE_LLM + return row + + +def _replacement_columns(selected_models: ReplaceModelSelection) -> list[ColumnConfigT]: + replace_alias = resolve_model_alias("replacement_generator", selected_models) + return [ + CustomColumnConfig( + name=COL_ENTITY_EXAMPLES, + generator_function=_prepare_replacement_inputs, + ), + LLMStructuredColumnConfig( + name=COL_REPLACEMENT_MAP_RAW, + prompt=_get_replacement_mapping_prompt(entities_column=COL_ENTITIES_FOR_REPLACE), + model_alias=replace_alias, + output_format=EntityReplacementMapSchema, + ), + CustomColumnConfig( + name=COL_REPLACEMENT_MAP, + generator_function=_filter_replacement_map, + ), + ] + + +def _canonical_row(row: dict[str, Any], mapping: dict[str, str]) -> dict[str, Any]: + canonical = row.copy() + for original, remapped in mapping.items(): + if remapped in row: + canonical[original] = row[remapped] + return canonical + + +def _copy_remapped_outputs( + row: dict[str, Any], + generated: dict[str, Any], + outputs: list[str], + mapping: dict[str, str], +) -> dict[str, Any]: + for output in outputs: + row[mapping.get(output, output)] = generated.get(output) + return row + + +def _remap_custom_column( + column: ColumnConfigT, + mapping: dict[str, str], + *, + skip: SkipConfig | None = None, +) -> CustomColumnConfig: + if not isinstance(column, CustomColumnConfig): + raise TypeError(f"Expected CustomColumnConfig, got {type(column).__name__}") + + generator = column.generator_function + metadata = generator.custom_column_metadata + required_columns = [mapping.get(name, name) for name in metadata["required_columns"]] + side_effect_columns = [mapping.get(name, name) for name in metadata["side_effect_columns"]] + model_aliases = list(metadata["model_aliases"]) + outputs = [column.name, *metadata["side_effect_columns"]] + + if model_aliases: + + @custom_column_generator( + required_columns=required_columns, + side_effect_columns=side_effect_columns, + model_aliases=model_aliases, + ) + def remapped_generator(row: dict[str, Any], generator_params: Any, models: dict) -> dict[str, Any]: + generated = generator(_canonical_row(row, mapping), generator_params, models) + return _copy_remapped_outputs(row, generated, outputs, mapping) + + elif column.generator_params is not None: + + @custom_column_generator( + required_columns=required_columns, + side_effect_columns=side_effect_columns, + ) + def remapped_generator(row: dict[str, Any], generator_params: Any) -> dict[str, Any]: + generated = generator(_canonical_row(row, mapping), generator_params) + return _copy_remapped_outputs(row, generated, outputs, mapping) + + else: + + @custom_column_generator( + required_columns=required_columns, + side_effect_columns=side_effect_columns, + ) + def remapped_generator(row: dict[str, Any]) -> dict[str, Any]: + generated = generator(_canonical_row(row, mapping)) + return _copy_remapped_outputs(row, generated, outputs, mapping) + + updates: dict[str, Any] = { + "name": mapping.get(column.name, column.name), + "generator_function": remapped_generator, + } + if skip is not None: + updates["skip"] = skip + return column.model_copy(update=updates) + + +def _evaluation_columns( + adapter: NddAdapter, + *, + selected_models: RewriteModelSelection, + evaluation: EvaluationCriteria, + state: EvaluationState, +) -> list[ColumnConfigT]: + mapping = { + COL_REWRITTEN_TEXT: state.rewritten_text, + COL_QUALITY_QA_REANSWER: state.quality_reanswer, + COL_PRIVACY_QA_REANSWER: state.privacy_reanswer, + COL_QUALITY_QA_COMPARE: state.quality_compare, + COL_UTILITY_SCORE: state.utility_score, + COL_LEAKAGE_MASS: state.leakage_mass, + COL_WEIGHTED_LEAKAGE_RATE: state.weighted_leakage_rate, + COL_ANY_HIGH_LEAKED: state.any_high_leaked, + COL_NEEDS_REPAIR: state.needs_repair, + } + return [ + _remap_custom_column(column, mapping) + for column in EvaluateWorkflow(adapter).columns( + selected_models=selected_models, + evaluation=evaluation, + ) + ] + + +def _repair_columns( + adapter: NddAdapter, + *, + selected_models: RewriteModelSelection, + privacy_goal: PrivacyGoal, + evaluation: EvaluationCriteria, + previous: EvaluationState, + state: RepairState, +) -> list[ColumnConfigT]: + mapping = { + COL_PRIVACY_QA_REANSWER: previous.privacy_reanswer, + COL_REWRITTEN_TEXT: previous.rewritten_text, + COL_LEAKAGE_MASS: previous.leakage_mass, + COL_WEIGHTED_LEAKAGE_RATE: previous.weighted_leakage_rate, + COL_ANY_HIGH_LEAKED: previous.any_high_leaked, + COL_UTILITY_SCORE: previous.utility_score, + COL_LEAKED_PRIVACY_ITEMS: state.leaked_items, + COL_REWRITTEN_TEXT_NEXT: state.rewritten_text, + } + columns = RepairWorkflow(adapter).columns( + selected_models=selected_models, + privacy_goal=privacy_goal, + effective_threshold=evaluation.repair_threshold, + ) + condition = SkipConfig(when=f"{{{{ not {previous.needs_repair} }}}}") + return [ + _remap_custom_column(column, mapping, skip=condition if index == 0 else None) + for index, column in enumerate(columns) + ] + + +def _finalization_column( + states: list[EvaluationState], + evaluation: EvaluationCriteria, +) -> CustomColumnConfig: + required_columns = list( + dict.fromkeys( + column + for state in states + for column in ( + state.rewritten_text, + state.quality_reanswer, + state.privacy_reanswer, + state.quality_compare, + state.utility_score, + state.leakage_mass, + state.weighted_leakage_rate, + state.any_high_leaked, + state.needs_repair, + ) + ) + ) + side_effect_columns = [ + COL_QUALITY_QA_REANSWER, + COL_PRIVACY_QA_REANSWER, + COL_QUALITY_QA_COMPARE, + COL_UTILITY_SCORE, + COL_LEAKAGE_MASS, + COL_WEIGHTED_LEAKAGE_RATE, + COL_ANY_HIGH_LEAKED, + COL_NEEDS_REPAIR, + COL_REPAIR_ITERATIONS, + COL_NEEDS_HUMAN_REVIEW, + ] + + @custom_column_generator( + required_columns=required_columns, + side_effect_columns=side_effect_columns, + ) + def finalize(row: dict[str, Any], generator_params: _FinalizationParams) -> dict[str, Any]: + state = next( + (candidate for candidate in reversed(states) if row.get(candidate.needs_repair) is not None), + states[0], + ) + row[COL_REWRITTEN_TEXT] = row.get(state.rewritten_text) + row[COL_QUALITY_QA_REANSWER] = normalize_payload(row.get(state.quality_reanswer)) + row[COL_PRIVACY_QA_REANSWER] = normalize_payload(row.get(state.privacy_reanswer)) + row[COL_QUALITY_QA_COMPARE] = normalize_payload(row.get(state.quality_compare)) + row[COL_UTILITY_SCORE] = row.get(state.utility_score) + row[COL_LEAKAGE_MASS] = row.get(state.leakage_mass) + row[COL_WEIGHTED_LEAKAGE_RATE] = row.get(state.weighted_leakage_rate) + row[COL_ANY_HIGH_LEAKED] = row.get(state.any_high_leaked) + row[COL_NEEDS_REPAIR] = row.get(state.needs_repair) + row[COL_REPAIR_ITERATIONS] = state.iteration + + needs_review = row[COL_REWRITTEN_TEXT] is None or bool(row[COL_ANY_HIGH_LEAKED]) + if generator_params.flag_utility_below is not None: + needs_review = needs_review or float(row[COL_UTILITY_SCORE]) < generator_params.flag_utility_below + if generator_params.flag_leakage_above is not None: + needs_review = needs_review or float(row[COL_LEAKAGE_MASS]) > generator_params.flag_leakage_above + row[COL_NEEDS_HUMAN_REVIEW] = needs_review + return row + + return CustomColumnConfig( + name=COL_REWRITTEN_TEXT, + generator_function=finalize, + generator_params=_FinalizationParams( + flag_utility_below=evaluation.flag_utility_below, + flag_leakage_above=evaluation.flag_leakage_above, + ), + propagate_skip=False, + ) + + +class CombinedRewriteWorkflow(RewriteWorkflow): + """Proof-of-concept rewrite workflow executed as one DataDesigner graph.""" + + def __init__(self, adapter: NddAdapter) -> None: + super().__init__(adapter) + + def build_graph( + self, + *, + selected_models: RewriteModelSelection, + replace_model_selection: ReplaceModelSelection, + privacy_goal: PrivacyGoal, + evaluation: EvaluationCriteria, + data_summary: str | None = None, + strict_entity_protection: bool = False, + ) -> CombinedRewriteGraph: + columns = _replacement_columns(replace_model_selection) + columns.extend( + DomainClassificationWorkflow().columns(selected_models=selected_models, data_summary=data_summary) + ) + columns.extend( + SensitivityDispositionWorkflow().columns( + selected_models=selected_models, + privacy_goal=privacy_goal, + data_summary=data_summary, + strict_entity_protection=strict_entity_protection, + ) + ) + columns.extend(QAGenerationWorkflow().columns(selected_models=selected_models)) + + rewrite_columns = RewriteGenerationWorkflow().columns( + selected_models=selected_models, + privacy_goal=privacy_goal, + data_summary=data_summary, + ) + columns.extend( + _remap_custom_column(column, {COL_REWRITTEN_TEXT: COL_REWRITTEN_TEXT_INITIAL}) + if column.name == COL_REWRITTEN_TEXT + else column + for column in rewrite_columns + ) + + evaluation_states = [_evaluation_state(0, COL_REWRITTEN_TEXT_INITIAL)] + repair_states: list[RepairState] = [] + columns.extend( + _evaluation_columns( + self._adapter, + selected_models=selected_models, + evaluation=evaluation, + state=evaluation_states[0], + ) + ) + + for iteration in range(evaluation.max_repair_iterations): + repair_state = _repair_state(iteration) + repair_states.append(repair_state) + columns.extend( + _repair_columns( + self._adapter, + selected_models=selected_models, + privacy_goal=privacy_goal, + evaluation=evaluation, + previous=evaluation_states[-1], + state=repair_state, + ) + ) + evaluation_state = _evaluation_state(iteration + 1, repair_state.rewritten_text) + evaluation_states.append(evaluation_state) + columns.extend( + _evaluation_columns( + self._adapter, + selected_models=selected_models, + evaluation=evaluation, + state=evaluation_state, + ) + ) + + columns.append(_finalization_column(evaluation_states, evaluation)) + internal_columns = [ + COL_ENTITY_EXAMPLES, + COL_ENTITIES_FOR_REPLACE, + COL_ENTITIES_FOR_REPLACE_JSON, + COL_REPLACEMENT_MAP_RAW, + COL_REWRITTEN_TEXT_INITIAL, + *(column for state in evaluation_states for column in state.__dict__.values() if isinstance(column, str)), + *(column for state in repair_states for column in state.__dict__.values() if isinstance(column, str)), + ] + return CombinedRewriteGraph( + columns=columns, + evaluation_states=evaluation_states, + repair_states=repair_states, + internal_columns=list(dict.fromkeys(internal_columns)), + ) + + def run( + self, + dataframe: pd.DataFrame, + *, + model_configs: list[ModelConfig], + selected_models: RewriteModelSelection, + replace_model_selection: ReplaceModelSelection, + privacy_goal: PrivacyGoal, + evaluation: EvaluationCriteria, + data_summary: str | None = None, + preview_num_records: int | None = None, + strict_entity_protection: bool = False, + ) -> RewriteResult: + with stage_timer("CombinedRewriteWorkflow.run", input_row_count=len(dataframe)) as measurement: + entity_rows, passthrough_rows = split_rows( + dataframe, + column=COL_ENTITIES_BY_VALUE, + predicate=_has_entities, + ) + measurement.update( + entity_row_count=len(entity_rows), + passthrough_row_count=len(passthrough_rows), + ) + if entity_rows.empty: + _apply_passthrough_defaults(passthrough_rows) + result = RewriteResult(dataframe=merge_and_reorder(passthrough_rows), failed_records=[]) + measurement.update(output_row_count=len(result.dataframe), failed_record_count=0) + return result + + graph = self.build_graph( + selected_models=selected_models, + replace_model_selection=replace_model_selection, + privacy_goal=privacy_goal, + evaluation=evaluation, + data_summary=data_summary, + strict_entity_protection=strict_entity_protection, + ) + seed = select_seed_cols(entity_rows, derive_seed_columns(graph.columns, entity_rows)) + run_result = self._adapter.run_workflow( + seed, + model_configs=model_configs, + columns=graph.columns, + workflow_name="rewrite-combined", + preview_num_records=preview_num_records, + ) + entity_rows = _join_new_columns(entity_rows, run_result.dataframe) + entity_rows = entity_rows.drop(columns=graph.internal_columns, errors="ignore") + _apply_passthrough_defaults(passthrough_rows) + result = RewriteResult( + dataframe=merge_and_reorder(entity_rows, passthrough_rows), + failed_records=run_result.failed_records, + ) + measurement.update( + output_row_count=len(result.dataframe), + failed_record_count=len(result.failed_records), + ) + return result diff --git a/src/anonymizer/engine/rewrite/rewrite_workflow.py b/src/anonymizer/engine/rewrite/rewrite_workflow.py index 8562302c..93940788 100644 --- a/src/anonymizer/engine/rewrite/rewrite_workflow.py +++ b/src/anonymizer/engine/rewrite/rewrite_workflow.py @@ -20,6 +20,9 @@ COL_LEAKAGE_MASS, COL_NEEDS_HUMAN_REVIEW, COL_NEEDS_REPAIR, + COL_PRIVACY_QA_REANSWER, + COL_QUALITY_QA_COMPARE, + COL_QUALITY_QA_REANSWER, COL_REPAIR_ITERATIONS, COL_REWRITTEN_TEXT, COL_REWRITTEN_TEXT_NEXT, @@ -54,6 +57,12 @@ COL_REPAIR_ITERATIONS: 0, } +_EVALUATION_PAYLOAD_COLUMNS = ( + COL_QUALITY_QA_REANSWER, + COL_PRIVACY_QA_REANSWER, + COL_QUALITY_QA_COMPARE, +) + def _detection_valid_fraction(row: pd.Series) -> float | None: """Convert bool COL_DETECTION_VALID to a 0–1 fraction for rewrite evaluate output. @@ -100,6 +109,12 @@ def _has_entities(entities_by_value: object) -> bool: return len(items) > 0 +def _normalize_evaluation_payloads(df: pd.DataFrame) -> None: + for column in _EVALUATION_PAYLOAD_COLUMNS: + if column in df.columns: + df[column] = df[column].map(normalize_payload) + + def _join_new_columns( target: pd.DataFrame, source: pd.DataFrame, @@ -366,6 +381,7 @@ def _run_evaluate_repair_loop( preview_num_records=preview_num_records, ) df = _join_new_columns(df, eval_result.dataframe, overwrite=True, seed_cols=eval_seed_cols) + _normalize_evaluation_payloads(df) all_failed.extend(eval_result.failed_records) repair_columns = self._repair_wf.columns( @@ -421,6 +437,7 @@ def _run_evaluate_repair_loop( failing_rows = _join_new_columns( failing_rows, eval_result.dataframe, overwrite=True, seed_cols=reeval_seed_cols ) + _normalize_evaluation_payloads(failing_rows) all_failed.extend(eval_result.failed_records) df = pd.concat([passing_rows, failing_rows], ignore_index=True) diff --git a/src/anonymizer/interface/anonymizer.py b/src/anonymizer/interface/anonymizer.py index d7c26eba..b1c60903 100644 --- a/src/anonymizer/interface/anonymizer.py +++ b/src/anonymizer/interface/anonymizer.py @@ -68,6 +68,7 @@ from anonymizer.engine.replace.llm_replace_workflow import LlmReplaceWorkflow from anonymizer.engine.replace.replace_runner import ReplacementWorkflow from anonymizer.engine.resolved_input import ResolvedInput +from anonymizer.engine.rewrite.combined_workflow import CombinedRewriteWorkflow from anonymizer.engine.rewrite.rewrite_workflow import RewriteWorkflow from anonymizer.engine.schemas import EntitiesByValueSchema from anonymizer.interface.errors import InvalidConfigError @@ -158,6 +159,7 @@ def __init__( detection_workflow: EntityDetectionWorkflow | None = None, replace_runner: ReplacementWorkflow | None = None, rewrite_runner: RewriteWorkflow | None = None, + combined_rewrite_runner: CombinedRewriteWorkflow | None = None, ) -> None: """Create an Anonymizer instance. @@ -178,6 +180,7 @@ def __init__( detection_workflow: Custom detection workflow (advanced/testing). replace_runner: Custom replacement workflow (advanced/testing). rewrite_runner: Custom rewrite workflow (advanced/testing). + combined_rewrite_runner: Custom combined rewrite workflow (advanced/testing). """ _initialize_logging() # Tag DataDesigner telemetry events so they're filterable as anonymizer traffic in @@ -227,6 +230,7 @@ def __init__( adapter=self._adapter, ) self._rewrite_runner = rewrite_runner or RewriteWorkflow(adapter=self._adapter) + self._combined_rewrite_runner = combined_rewrite_runner or CombinedRewriteWorkflow(adapter=self._adapter) def run( self, @@ -757,7 +761,10 @@ def _run_internal_impl( privacy_goal = config.rewrite.privacy_goal if privacy_goal is None: raise InvalidConfigError("rewrite.privacy_goal must not be None") - rewrite_result = self._rewrite_runner.run( + rewrite_runner = ( + self._combined_rewrite_runner if config.rewrite.use_combined_graph else self._rewrite_runner + ) + rewrite_result = rewrite_runner.run( detection_result.dataframe, model_configs=self._model_configs, selected_models=self._selected_models.rewrite, diff --git a/tests/config/test_rewrite.py b/tests/config/test_rewrite.py index c98e383e..c1426114 100644 --- a/tests/config/test_rewrite.py +++ b/tests/config/test_rewrite.py @@ -5,6 +5,7 @@ import pytest +from anonymizer.config.anonymizer_config import Rewrite from anonymizer.config.rewrite import ( EvaluationCriteria, PrivacyGoal, @@ -75,6 +76,10 @@ def test_default_is_low() -> None: assert criteria.max_repair_iterations == 3 +def test_combined_graph_is_opt_in() -> None: + assert Rewrite().use_combined_graph is False + + def test_minimal_bundles_aggressive_review_flags() -> None: criteria = EvaluationCriteria(risk_tolerance=RiskTolerance.minimal) assert criteria.flag_utility_below == 0.6 diff --git a/tests/engine/test_combined_rewrite_workflow.py b/tests/engine/test_combined_rewrite_workflow.py new file mode 100644 index 00000000..5b43946a --- /dev/null +++ b/tests/engine/test_combined_rewrite_workflow.py @@ -0,0 +1,606 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from pathlib import Path +from typing import Any +from unittest.mock import Mock, patch + +import numpy as np +import pandas as pd +import pytest +from data_designer.config import SkipConfig, custom_column_generator +from data_designer.config.column_configs import CustomColumnConfig +from data_designer.config.column_types import ColumnConfigT +from data_designer.config.models import ModelConfig, ModelProvider +from data_designer.interface.data_designer import DataDesigner + +from anonymizer.config.models import ModelSelection, ReplaceModelSelection, RewriteModelSelection +from anonymizer.config.rewrite import EvaluationCriteria, PrivacyGoal +from anonymizer.engine.constants import ( + COL_ANY_HIGH_LEAKED, + COL_ENTITIES_BY_VALUE, + COL_FULL_REWRITE, + COL_LATENT_ENTITIES, + COL_LEAKAGE_MASS, + COL_NEEDS_HUMAN_REVIEW, + COL_NEEDS_REPAIR, + COL_REPAIR_ITERATIONS, + COL_REWRITTEN_TEXT, + COL_REWRITTEN_TEXT_INITIAL, + COL_REWRITTEN_TEXT_NEXT, + COL_TAG_NOTATION, + COL_TAGGED_TEXT, + COL_TEXT, + COL_UTILITY_SCORE, + COL_WEIGHTED_LEAKAGE_RATE, +) +from anonymizer.engine.ndd.adapter import RECORD_ID_COLUMN, FailedRecord, NddAdapter, WorkflowRunResult +from anonymizer.engine.rewrite.combined_workflow import ( + CombinedRewriteGraph, + CombinedRewriteWorkflow, + EvaluationState, + RepairState, +) +from anonymizer.engine.rewrite.rewrite_workflow import RewriteWorkflow +from anonymizer.measurement import MeasurementCollector, measurement_session + +_PRIVACY_GOAL = PrivacyGoal( + protect="All direct identifiers including names, locations, and contact details", + preserve="Career trajectory, skills, and professional context in abstract terms", +) +_REPAIRS_NEEDED = "repairs_needed" +_REPLACE_PATCH = "anonymizer.engine.rewrite.rewrite_workflow.LlmReplaceWorkflow" + + +@custom_column_generator(required_columns=[_REPAIRS_NEEDED]) +def _initial_rewrite(row: dict[str, Any]) -> dict[str, Any]: + row[COL_REWRITTEN_TEXT_INITIAL] = "rewrite-0" + return row + + +def _deterministic_evaluation_column(state: EvaluationState) -> CustomColumnConfig: + side_effect_columns = [ + state.privacy_reanswer, + state.quality_compare, + state.utility_score, + state.leakage_mass, + state.weighted_leakage_rate, + state.any_high_leaked, + state.needs_repair, + ] + + @custom_column_generator( + required_columns=[_REPAIRS_NEEDED, state.rewritten_text], + side_effect_columns=side_effect_columns, + ) + def evaluate(row: dict[str, Any]) -> dict[str, Any]: + needs_repair = int(row[_REPAIRS_NEEDED]) > state.iteration + items: list[Any] | np.ndarray = [] if state.iteration == 0 else np.array([], dtype=object) + row[state.quality_reanswer] = {"answers": items} + row[state.privacy_reanswer] = {"answers": items} + row[state.quality_compare] = {"per_item": items} + row[state.utility_score] = 0.5 if needs_repair else 1.0 + row[state.leakage_mass] = 1.0 if needs_repair else 0.0 + row[state.weighted_leakage_rate] = 1.0 if needs_repair else 0.0 + row[state.any_high_leaked] = needs_repair + row[state.needs_repair] = needs_repair + return row + + return CustomColumnConfig(name=state.quality_reanswer, generator_function=evaluate) + + +def _deterministic_repair_column(previous: EvaluationState, state: RepairState) -> CustomColumnConfig: + @custom_column_generator(required_columns=[previous.needs_repair]) + def repair(row: dict[str, Any]) -> dict[str, Any]: + row[state.rewritten_text] = f"rewrite-{state.iteration + 1}" + return row + + return CustomColumnConfig( + name=state.rewritten_text, + generator_function=repair, + skip=SkipConfig(when=f"{{{{ not {previous.needs_repair} }}}}"), + ) + + +def _deterministic_columns(graph: CombinedRewriteGraph) -> list[ColumnConfigT]: + columns: list[ColumnConfigT] = [ + CustomColumnConfig( + name=graph.evaluation_states[0].rewritten_text, + generator_function=_initial_rewrite, + ), + _deterministic_evaluation_column(graph.evaluation_states[0]), + ] + for previous, repair, current in zip( + graph.evaluation_states, + graph.repair_states, + graph.evaluation_states[1:], + ): + columns.extend( + [ + _deterministic_repair_column(previous, repair), + _deterministic_evaluation_column(current), + ] + ) + columns.append(graph.columns[-1]) + return columns + + +@pytest.mark.parametrize("max_repair_iterations", [0, 1, 3]) +def test_graph_unrolls_conditional_repairs( + stub_rewrite_model_selection: RewriteModelSelection, + stub_replace_model_selection: ReplaceModelSelection, + max_repair_iterations: int, +) -> None: + graph = CombinedRewriteWorkflow(adapter=Mock()).build_graph( + selected_models=stub_rewrite_model_selection, + replace_model_selection=stub_replace_model_selection, + privacy_goal=_PRIVACY_GOAL, + evaluation=EvaluationCriteria(max_repair_iterations=max_repair_iterations), + ) + + assert len(graph.evaluation_states) == max_repair_iterations + 1 + assert len(graph.repair_states) == max_repair_iterations + by_name = {column.name: column for column in graph.columns} + for previous, repair in zip(graph.evaluation_states, graph.repair_states): + condition = by_name[repair.leaked_items].skip + assert condition is not None + assert previous.needs_repair in condition.columns + + +def test_finalizer_selects_last_executed_iteration( + stub_rewrite_model_selection: RewriteModelSelection, + stub_replace_model_selection: ReplaceModelSelection, +) -> None: + graph = CombinedRewriteWorkflow(adapter=Mock()).build_graph( + selected_models=stub_rewrite_model_selection, + replace_model_selection=stub_replace_model_selection, + privacy_goal=_PRIVACY_GOAL, + evaluation=EvaluationCriteria(max_repair_iterations=2), + ) + initial, repaired, skipped = graph.evaluation_states + row = { + initial.rewritten_text: "Initial rewrite", + initial.quality_reanswer: {"answers": []}, + initial.privacy_reanswer: {"answers": []}, + initial.quality_compare: {"per_item": []}, + initial.utility_score: 0.7, + initial.leakage_mass: 2.0, + initial.weighted_leakage_rate: 0.8, + initial.any_high_leaked: True, + initial.needs_repair: True, + repaired.rewritten_text: "Repaired rewrite", + repaired.quality_reanswer: {"answers": []}, + repaired.privacy_reanswer: {"answers": []}, + repaired.quality_compare: {"per_item": []}, + repaired.utility_score: 0.9, + repaired.leakage_mass: 0.1, + repaired.weighted_leakage_rate: 0.05, + repaired.any_high_leaked: False, + repaired.needs_repair: False, + skipped.needs_repair: None, + } + finalizer = graph.columns[-1] + assert isinstance(finalizer, CustomColumnConfig) + + result = finalizer.generator_function(row, finalizer.generator_params) + + assert result[COL_REWRITTEN_TEXT] == "Repaired rewrite" + assert result[COL_UTILITY_SCORE] == 0.9 + assert result[COL_LEAKAGE_MASS] == 0.1 + assert result[COL_WEIGHTED_LEAKAGE_RATE] == 0.05 + assert result[COL_ANY_HIGH_LEAKED] is False + assert result[COL_NEEDS_REPAIR] is False + assert result[COL_REPAIR_ITERATIONS] == 1 + assert result[COL_NEEDS_HUMAN_REVIEW] is False + + +def test_combined_graph_preserves_malformed_rewrite_handling( + stub_rewrite_model_selection: RewriteModelSelection, + stub_replace_model_selection: ReplaceModelSelection, +) -> None: + graph = CombinedRewriteWorkflow(adapter=Mock()).build_graph( + selected_models=stub_rewrite_model_selection, + replace_model_selection=stub_replace_model_selection, + privacy_goal=_PRIVACY_GOAL, + evaluation=EvaluationCriteria(max_repair_iterations=0), + ) + rewrite_column = next(column for column in graph.columns if column.name == COL_REWRITTEN_TEXT_INITIAL) + assert isinstance(rewrite_column, CustomColumnConfig) + + result = rewrite_column.generator_function({COL_FULL_REWRITE: "not-a-valid-payload"}) + + assert result[COL_REWRITTEN_TEXT_INITIAL] is None + + +def test_conditional_repairs_execute_independently_per_row( + tmp_path: Path, + stub_rewrite_model_selection: RewriteModelSelection, + stub_replace_model_selection: ReplaceModelSelection, +) -> None: + data_designer = DataDesigner(artifact_path=tmp_path / "artifacts", auto_configure_logging=False) + adapter = NddAdapter(data_designer) + graph = CombinedRewriteWorkflow(adapter).build_graph( + selected_models=stub_rewrite_model_selection, + replace_model_selection=stub_replace_model_selection, + privacy_goal=_PRIVACY_GOAL, + evaluation=EvaluationCriteria(max_repair_iterations=2), + ) + columns = _deterministic_columns(graph) + collector = MeasurementCollector(record_hash_key="test-key") + + with measurement_session(collector): + result = adapter.run_workflow( + pd.DataFrame({_REPAIRS_NEEDED: [0, 1, 2, 3]}), + model_configs=[], + columns=columns, + workflow_name="rewrite-combined", + preview_num_records=4, + ) + + assert result.failed_records == [] + assert result.dataframe[COL_REWRITTEN_TEXT].tolist() == [ + "rewrite-0", + "rewrite-1", + "rewrite-2", + "rewrite-2", + ] + assert result.dataframe[COL_REPAIR_ITERATIONS].tolist() == [0, 1, 2, 2] + assert result.dataframe[COL_NEEDS_REPAIR].tolist() == [False, False, False, True] + assert result.dataframe[COL_NEEDS_HUMAN_REVIEW].tolist() == [False, False, False, True] + workflow_records = [record for record in collector.records if record["record_type"] == "ndd_workflow"] + assert len(workflow_records) == 1 + assert workflow_records[0]["workflow_name"] == "rewrite-combined" + assert workflow_records[0]["input_row_count"] == 4 + assert workflow_records[0]["output_row_count"] == 4 + assert workflow_records[0]["column_count"] == len(columns) + + +def test_run_executes_one_data_designer_workflow( + stub_model_configs: list[ModelConfig], + stub_rewrite_model_selection: RewriteModelSelection, + stub_replace_model_selection: ReplaceModelSelection, +) -> None: + dataframe = pd.DataFrame( + { + COL_TEXT: ["Alice works at Acme"], + COL_ENTITIES_BY_VALUE: [{"entities_by_value": [{"value": "Alice", "labels": ["first_name"]}]}], + } + ) + output = dataframe.copy() + output[COL_REWRITTEN_TEXT] = "Maria works at a company" + output[COL_UTILITY_SCORE] = 0.9 + output[COL_LEAKAGE_MASS] = 0.1 + output[COL_WEIGHTED_LEAKAGE_RATE] = 0.05 + output[COL_ANY_HIGH_LEAKED] = False + output[COL_NEEDS_REPAIR] = False + output[COL_REPAIR_ITERATIONS] = 1 + output[COL_NEEDS_HUMAN_REVIEW] = False + adapter = Mock() + adapter.run_workflow.return_value = WorkflowRunResult(dataframe=output, failed_records=[]) + + collector = MeasurementCollector(record_hash_key="test-key") + with measurement_session(collector): + result = CombinedRewriteWorkflow(adapter=adapter).run( + dataframe, + model_configs=stub_model_configs, + selected_models=stub_rewrite_model_selection, + replace_model_selection=stub_replace_model_selection, + privacy_goal=_PRIVACY_GOAL, + evaluation=EvaluationCriteria(max_repair_iterations=2), + ) + + adapter.run_workflow.assert_called_once() + assert adapter.run_workflow.call_args.kwargs["workflow_name"] == "rewrite-combined" + assert result.dataframe[COL_REWRITTEN_TEXT].tolist() == ["Maria works at a company"] + stage_records = [record for record in collector.records if record["record_type"] == "stage"] + assert len(stage_records) == 1 + assert stage_records[0]["stage"] == "CombinedRewriteWorkflow.run" + assert stage_records[0]["input_row_count"] == 1 + assert stage_records[0]["output_row_count"] == 1 + assert stage_records[0]["failed_record_count"] == 0 + + +def test_run_reports_combined_failure_and_drops_only_failed_entity_row( + stub_model_configs: list[ModelConfig], + stub_rewrite_model_selection: RewriteModelSelection, + stub_replace_model_selection: ReplaceModelSelection, +) -> None: + dataframe = pd.DataFrame( + { + RECORD_ID_COLUMN: ["entity-ok", "passthrough", "entity-failed"], + COL_TEXT: ["Alice works at Acme", "No entities", "Bob works at Beta"], + COL_ENTITIES_BY_VALUE: [ + {"entities_by_value": [{"value": "Alice", "labels": ["first_name"]}]}, + {"entities_by_value": []}, + {"entities_by_value": [{"value": "Bob", "labels": ["first_name"]}]}, + ], + } + ) + output = dataframe.iloc[[0]].copy() + output[COL_REWRITTEN_TEXT] = "Person works at Company" + output[COL_UTILITY_SCORE] = 0.9 + output[COL_LEAKAGE_MASS] = 0.0 + output[COL_WEIGHTED_LEAKAGE_RATE] = 0.0 + output[COL_ANY_HIGH_LEAKED] = False + output[COL_NEEDS_REPAIR] = False + output[COL_REPAIR_ITERATIONS] = 0 + output[COL_NEEDS_HUMAN_REVIEW] = False + failed = FailedRecord( + record_id="entity-failed", + step="rewrite-combined", + reason="Record missing from workflow output", + ) + adapter = Mock() + adapter.run_workflow.return_value = WorkflowRunResult(dataframe=output, failed_records=[failed]) + + result = CombinedRewriteWorkflow(adapter=adapter).run( + dataframe, + model_configs=stub_model_configs, + selected_models=stub_rewrite_model_selection, + replace_model_selection=stub_replace_model_selection, + privacy_goal=_PRIVACY_GOAL, + evaluation=EvaluationCriteria(max_repair_iterations=2), + ) + + assert result.dataframe[RECORD_ID_COLUMN].tolist() == ["entity-ok", "passthrough"] + assert result.dataframe[COL_REWRITTEN_TEXT].tolist() == ["Person works at Company", "No entities"] + assert result.failed_records == [failed] + + +def test_combined_and_legacy_paths_return_equivalent_repaired_result( + stub_model_configs: list[ModelConfig], + stub_rewrite_model_selection: RewriteModelSelection, + stub_replace_model_selection: ReplaceModelSelection, +) -> None: + dataframe = pd.DataFrame( + { + COL_TEXT: ["Alice works at Acme"], + COL_ENTITIES_BY_VALUE: [{"entities_by_value": [{"value": "Alice", "labels": ["first_name"]}]}], + } + ) + pipeline = dataframe.copy() + pipeline[COL_REWRITTEN_TEXT] = "Initial rewrite" + evaluation_before = pd.DataFrame( + { + COL_NEEDS_REPAIR: [True], + COL_UTILITY_SCORE: [0.5], + COL_LEAKAGE_MASS: [1.0], + COL_WEIGHTED_LEAKAGE_RATE: [1.0], + COL_ANY_HIGH_LEAKED: [True], + } + ) + repair = pd.DataFrame({COL_REWRITTEN_TEXT_NEXT: ["Repaired rewrite"]}) + evaluation_after = pd.DataFrame( + { + COL_NEEDS_REPAIR: [False], + COL_UTILITY_SCORE: [0.9], + COL_LEAKAGE_MASS: [0.0], + COL_WEIGHTED_LEAKAGE_RATE: [0.0], + COL_ANY_HIGH_LEAKED: [False], + } + ) + legacy_adapter = Mock() + legacy_adapter.run_workflow.side_effect = [ + WorkflowRunResult(dataframe=pipeline, failed_records=[]), + WorkflowRunResult(dataframe=evaluation_before, failed_records=[]), + WorkflowRunResult(dataframe=repair, failed_records=[]), + WorkflowRunResult(dataframe=evaluation_after, failed_records=[]), + ] + with patch(_REPLACE_PATCH) as replace_workflow: + replace_workflow.return_value.generate_map_only.return_value = WorkflowRunResult( + dataframe=dataframe.copy(), + failed_records=[], + ) + legacy_result = RewriteWorkflow(adapter=legacy_adapter).run( + dataframe, + model_configs=stub_model_configs, + selected_models=stub_rewrite_model_selection, + replace_model_selection=stub_replace_model_selection, + privacy_goal=_PRIVACY_GOAL, + evaluation=EvaluationCriteria(max_repair_iterations=2), + ) + + combined_output = dataframe.copy() + combined_output[COL_REWRITTEN_TEXT] = "Repaired rewrite" + combined_output[COL_UTILITY_SCORE] = 0.9 + combined_output[COL_LEAKAGE_MASS] = 0.0 + combined_output[COL_WEIGHTED_LEAKAGE_RATE] = 0.0 + combined_output[COL_ANY_HIGH_LEAKED] = False + combined_output[COL_NEEDS_REPAIR] = False + combined_output[COL_REPAIR_ITERATIONS] = 1 + combined_output[COL_NEEDS_HUMAN_REVIEW] = False + combined_adapter = Mock() + combined_adapter.run_workflow.return_value = WorkflowRunResult( + dataframe=combined_output, + failed_records=[], + ) + combined_result = CombinedRewriteWorkflow(adapter=combined_adapter).run( + dataframe, + model_configs=stub_model_configs, + selected_models=stub_rewrite_model_selection, + replace_model_selection=stub_replace_model_selection, + privacy_goal=_PRIVACY_GOAL, + evaluation=EvaluationCriteria(max_repair_iterations=2), + ) + + output_columns = [ + COL_REWRITTEN_TEXT, + COL_UTILITY_SCORE, + COL_LEAKAGE_MASS, + COL_WEIGHTED_LEAKAGE_RATE, + COL_ANY_HIGH_LEAKED, + COL_NEEDS_REPAIR, + COL_REPAIR_ITERATIONS, + COL_NEEDS_HUMAN_REVIEW, + ] + pd.testing.assert_frame_equal( + legacy_result.dataframe[output_columns], + combined_result.dataframe[output_columns], + check_dtype=False, + ) + assert legacy_result.failed_records == combined_result.failed_records == [] + + +@pytest.mark.parametrize( + ("repairs_needed", "expected_counts"), + [ + ([0] * 62 + [1, 2], {0: 62, 1: 1, 2: 1}), + ([2] * 62 + [0, 1], {0: 1, 1: 1, 2: 62}), + ], + ids=["mostly-skipped", "mostly-repaired"], +) +def test_conditional_graph_handles_larger_mixed_batches( + tmp_path: Path, + stub_rewrite_model_selection: RewriteModelSelection, + stub_replace_model_selection: ReplaceModelSelection, + repairs_needed: list[int], + expected_counts: dict[int, int], +) -> None: + data_designer = DataDesigner(artifact_path=tmp_path / "artifacts", auto_configure_logging=False) + adapter = NddAdapter(data_designer) + graph = CombinedRewriteWorkflow(adapter).build_graph( + selected_models=stub_rewrite_model_selection, + replace_model_selection=stub_replace_model_selection, + privacy_goal=_PRIVACY_GOAL, + evaluation=EvaluationCriteria(max_repair_iterations=2), + ) + + result = adapter.run_workflow( + pd.DataFrame({_REPAIRS_NEEDED: repairs_needed}), + model_configs=[], + columns=_deterministic_columns(graph), + workflow_name="rewrite-combined-scale-test", + preview_num_records=len(repairs_needed), + ) + + assert result.failed_records == [] + assert result.dataframe[COL_REPAIR_ITERATIONS].value_counts().to_dict() == expected_counts + assert len(result.dataframe) == len(repairs_needed) + + +def test_run_preserves_mixed_row_order_and_passthrough_defaults( + stub_model_configs: list[ModelConfig], + stub_rewrite_model_selection: RewriteModelSelection, + stub_replace_model_selection: ReplaceModelSelection, +) -> None: + dataframe = pd.DataFrame( + { + COL_TEXT: ["Alice works at Acme", "No entities", "Bob works at Beta"], + COL_ENTITIES_BY_VALUE: [ + {"entities_by_value": [{"value": "Alice", "labels": ["first_name"]}]}, + {"entities_by_value": []}, + {"entities_by_value": [{"value": "Bob", "labels": ["first_name"]}]}, + ], + } + ) + adapter = Mock() + + def run_workflow(entity_rows: pd.DataFrame, **_: Any) -> WorkflowRunResult: + output = entity_rows.copy() + output[COL_REWRITTEN_TEXT] = ["Person works at Company", "Worker works at Business"] + output[COL_UTILITY_SCORE] = 0.9 + output[COL_LEAKAGE_MASS] = 0.0 + output[COL_WEIGHTED_LEAKAGE_RATE] = 0.0 + output[COL_ANY_HIGH_LEAKED] = False + output[COL_NEEDS_REPAIR] = False + output[COL_REPAIR_ITERATIONS] = 0 + output[COL_NEEDS_HUMAN_REVIEW] = False + return WorkflowRunResult(dataframe=output, failed_records=[]) + + adapter.run_workflow.side_effect = run_workflow + + result = CombinedRewriteWorkflow(adapter=adapter).run( + dataframe, + model_configs=stub_model_configs, + selected_models=stub_rewrite_model_selection, + replace_model_selection=stub_replace_model_selection, + privacy_goal=_PRIVACY_GOAL, + evaluation=EvaluationCriteria(max_repair_iterations=2), + ) + + assert result.dataframe[COL_TEXT].tolist() == dataframe[COL_TEXT].tolist() + assert result.dataframe[COL_REWRITTEN_TEXT].tolist() == [ + "Person works at Company", + "No entities", + "Worker works at Business", + ] + assert result.dataframe[COL_UTILITY_SCORE].tolist() == [0.9, 1.0, 0.9] + assert result.dataframe[COL_REPAIR_ITERATIONS].tolist() == [0, 0, 0] + assert adapter.run_workflow.call_args.args[0][COL_TEXT].tolist() == [ + "Alice works at Acme", + "Bob works at Beta", + ] + + +def test_run_skips_data_designer_when_no_rows_have_entities( + stub_model_configs: list[ModelConfig], + stub_rewrite_model_selection: RewriteModelSelection, + stub_replace_model_selection: ReplaceModelSelection, +) -> None: + dataframe = pd.DataFrame( + { + COL_TEXT: ["No entities", "Still no entities"], + COL_ENTITIES_BY_VALUE: [{"entities_by_value": []}, {"entities_by_value": []}], + } + ) + adapter = Mock() + + result = CombinedRewriteWorkflow(adapter=adapter).run( + dataframe, + model_configs=stub_model_configs, + selected_models=stub_rewrite_model_selection, + replace_model_selection=stub_replace_model_selection, + privacy_goal=_PRIVACY_GOAL, + evaluation=EvaluationCriteria(max_repair_iterations=2), + ) + + adapter.run_workflow.assert_not_called() + assert result.dataframe[COL_REWRITTEN_TEXT].tolist() == dataframe[COL_TEXT].tolist() + assert result.dataframe[COL_UTILITY_SCORE].tolist() == [1.0, 1.0] + assert result.dataframe[COL_REPAIR_ITERATIONS].tolist() == [0, 0] + + +@pytest.mark.parametrize("max_repair_iterations", [0, 2, 10]) +def test_graph_validates_with_data_designer( + tmp_path: Path, + stub_slim_model_selection: ModelSelection, + max_repair_iterations: int, +) -> None: + provider = ModelProvider( + name="stub", + endpoint="http://stub.invalid/v1", + provider_type="openai", + api_key="EMPTY", + ) + data_designer = DataDesigner( + artifact_path=tmp_path / "artifacts", + model_providers=[provider], + auto_configure_logging=False, + ) + adapter = NddAdapter(data_designer) + graph = CombinedRewriteWorkflow(adapter).build_graph( + selected_models=stub_slim_model_selection.rewrite, + replace_model_selection=stub_slim_model_selection.replace, + privacy_goal=_PRIVACY_GOAL, + evaluation=EvaluationCriteria(max_repair_iterations=max_repair_iterations), + ) + dataframe = pd.DataFrame( + { + COL_TEXT: ["Alice works at Acme"], + COL_TAGGED_TEXT: ["[[Alice|first_name]] works at Acme"], + COL_TAG_NOTATION: ["bracket"], + COL_ENTITIES_BY_VALUE: [{"entities_by_value": [{"value": "Alice", "labels": ["first_name"]}]}], + COL_LATENT_ENTITIES: [{"latent_entities": []}], + } + ) + model_configs = [ModelConfig(alias="known", model="stub-model", provider="stub")] + builder = adapter.build_config( + dataframe, + model_configs=model_configs, + columns=graph.columns, + seed_path=tmp_path / "seed.parquet", + ) + + data_designer.validate(builder) diff --git a/tests/engine/test_rewrite_workflow.py b/tests/engine/test_rewrite_workflow.py index 6e7a824c..609378d2 100644 --- a/tests/engine/test_rewrite_workflow.py +++ b/tests/engine/test_rewrite_workflow.py @@ -3,8 +3,10 @@ from __future__ import annotations +import json from unittest.mock import Mock, patch +import numpy as np import pandas as pd import pytest from data_designer.config.models import ModelConfig @@ -21,6 +23,7 @@ COL_LEAKAGE_MASS, COL_NEEDS_HUMAN_REVIEW, COL_NEEDS_REPAIR, + COL_PRIVACY_QA_REANSWER, COL_REPAIR_ITERATIONS, COL_REWRITTEN_TEXT, COL_REWRITTEN_TEXT_NEXT, @@ -567,6 +570,10 @@ def test_only_failing_rows_sent_to_repair( eval_df[COL_UTILITY_SCORE] = [0.9, 0.9] eval_df[COL_LEAKAGE_MASS] = [2.0, 0.1] eval_df[COL_ANY_HIGH_LEAKED] = [True, False] + eval_df[COL_PRIVACY_QA_REANSWER] = [ + {"answers": np.array([{"id": 0}], dtype=object)}, + {"answers": []}, + ] failing_row = eval_df[eval_df[COL_NEEDS_REPAIR]].copy() repaired_row = failing_row.copy() @@ -606,6 +613,7 @@ def test_only_failing_rows_sent_to_repair( assert len(repair_calls) == 1 repair_input_df = repair_calls[0].args[0] assert len(repair_input_df) == 1 + assert json.loads(repair_input_df[COL_PRIVACY_QA_REANSWER].iloc[0]) == {"answers": [{"id": 0}]} def test_repair_iterations_tracked_per_row( diff --git a/tests/interface/test_anonymizer_interface.py b/tests/interface/test_anonymizer_interface.py index ba1bea59..797e0127 100644 --- a/tests/interface/test_anonymizer_interface.py +++ b/tests/interface/test_anonymizer_interface.py @@ -721,6 +721,19 @@ def test_run_rewrite_calls_rewrite_runner(stub_input: AnonymizerInput) -> None: assert call_kwargs["evaluation"] == config.rewrite.evaluation +def test_run_rewrite_uses_combined_runner_when_enabled(stub_input: AnonymizerInput) -> None: + config = AnonymizerConfig(rewrite=Rewrite(use_combined_graph=True)) + anonymizer, _, _, rewrite_runner = _make_anonymizer() + combined_runner = Mock(spec=RewriteWorkflow) + combined_runner.run.return_value = rewrite_runner.run.return_value + anonymizer._combined_rewrite_runner = combined_runner + + anonymizer.run(config=config, data=stub_input) + + rewrite_runner.run.assert_not_called() + combined_runner.run.assert_called_once() + + def test_run_rewrite_rejects_missing_privacy_goal_before_workflows(stub_input: AnonymizerInput) -> None: config = AnonymizerConfig(rewrite=Rewrite()) assert config.rewrite is not None