diff --git a/docs/in_memory_api.md b/docs/in_memory_api.md index 9dc4fd4..146b052 100644 --- a/docs/in_memory_api.md +++ b/docs/in_memory_api.md @@ -127,7 +127,8 @@ python3 bin/main.py \ ### Supported CLI Flags * `--dataset`: Path to the input CSV file. (Supports standard CSV parsing - arguments via `--read_csv_args`). + arguments via `--field_separator`, `--column_names`, and + `--column_count`). * `--domain`: Path to the YAML domain specification file. * `--epsilon`, `--delta`: Total DP privacy budget. * `--mechanism`: Supported options are `mst`, `aim`, `independent`, and diff --git a/dpsynth/bin/comparison.py b/dpsynth/bin/comparison.py index f6cd85d..eda4c2c 100644 --- a/dpsynth/bin/comparison.py +++ b/dpsynth/bin/comparison.py @@ -22,16 +22,12 @@ from absl import flags from absl import logging from dpsynth import domain -import fancyflags as ff import pandas as pd -import sdmetrics import pathlib PathType = pathlib.Path -QualityReport = sdmetrics.reports.single_table.QualityReport -DiagnosticsReport = sdmetrics.reports.single_table.DiagnosticReport @dataclasses.dataclass(frozen=True) @@ -84,19 +80,16 @@ def __post_init__(self): ) -_COLUMNS_TO_CROSS_COMPARE = ff.DEFINE_dict( - 'cross_compare', - categorical_columns=ff.MultiString( - ['sex', 'race'], - 'Multiple categorical columns to use for cross-column comparison. We' - ' will group by all the categorical columns and compare the mean with' - ' each of the numerical columns.', - ), - numerical_columns=ff.MultiString( - ['age'], - 'Multiple numerical columns to cross compare, we will compare the mean' - ' of each column for all the categorical columns.', - ), +_CROSS_COMPARE_CATEGORICAL_COLUMNS = flags.DEFINE_list( + 'cross_compare_categorical_columns', + ['sex', 'race'], + 'Categorical columns to group by for cross-column comparison.', +) + +_CROSS_COMPARE_NUMERICAL_COLUMNS = flags.DEFINE_list( + 'cross_compare_numerical_columns', + ['age'], + 'Numerical columns whose grouped means are compared.', ) @@ -218,19 +211,30 @@ def _create_metadata_from_domain_yaml( elif isinstance(attr, domain.CategoricalAttribute): cols[name] = {'sdtype': 'categorical'} else: - raise ValueError('Unknown attribute type: {type(attr)}') + raise ValueError(f'Unknown attribute type: {type(attr)}') metadata['columns'] = cols return metadata +def _get_sdmetrics_report_classes() -> tuple[type[Any], type[Any]]: + try: + from sdmetrics.reports import single_table + except ImportError as exc: + raise ImportError( + 'comparison.py requires sdmetrics to generate quality reports.' + ) from exc + return single_table.QualityReport, single_table.DiagnosticReport + + def _create_quality_report( real_data: pd.DataFrame, synthetic_data: pd.DataFrame, metadata: dict[str, Any], - ) -> QualityReport: +) -> Any: """Creates a quality report.""" print('Quality report:') + QualityReport, _ = _get_sdmetrics_report_classes() quality_report = QualityReport() quality_report.generate(real_data, synthetic_data, metadata) if not quality_report.is_generated: @@ -243,9 +247,10 @@ def _create_diagnostics_report( real_data: pd.DataFrame, synthetic_data: pd.DataFrame, metadata: dict[str, Any], - ) -> DiagnosticsReport: +) -> Any: """Creates a diagnostics report.""" print('Diagnostics report:') + _, DiagnosticsReport = _get_sdmetrics_report_classes() diagnostics_report = DiagnosticsReport() diagnostics_report.generate(real_data, synthetic_data, metadata) if not diagnostics_report.is_generated: @@ -273,7 +278,8 @@ def main(_) -> None: sdmetrics_metadata = _create_metadata_from_domain_yaml(domain_path) columns_to_compare = _COLUMNS_TO_COMPARE.value columns_to_cross_compare = CompareGroupByColumns( - **_COLUMNS_TO_CROSS_COMPARE.value + categorical_columns=_CROSS_COMPARE_CATEGORICAL_COLUMNS.value or [], + numerical_columns=_CROSS_COMPARE_NUMERICAL_COLUMNS.value or [], ) # Compare histograms of the given columns. diff --git a/dpsynth/bin/main.py b/dpsynth/bin/main.py index 405f237..fd56243 100644 --- a/dpsynth/bin/main.py +++ b/dpsynth/bin/main.py @@ -28,7 +28,6 @@ from absl import flags import dpsynth from dpsynth.bin import _read_csv_args -import fancyflags as ff import numpy as np import pandas as pd @@ -80,16 +79,33 @@ required=True, ) -_READ_CSV_ARGS = ff.DEFINE_auto( - 'read_csv_args', - _read_csv_args.ReadCsvArgs, - _read_csv_args.FLAG_HELP, +_CSV_FIELD_SEPARATOR = flags.DEFINE_enum( + 'field_separator', + None, + ['tab', 'pipe', 'comma', 'semicolon', 'space'], + 'Field separator for the input CSV.', +) + +_CSV_COLUMN_NAMES = flags.DEFINE_list( + 'column_names', + None, + 'Column names to use when the input CSV does not have a header.', +) + +_CSV_COLUMN_COUNT = flags.DEFINE_integer( + 'column_count', + None, + 'Number of columns when the input CSV does not have a header.', ) def main(_): np.random.seed(_SEED.value) - read_csv_kwargs = _READ_CSV_ARGS.value().to_read_csv_kwargs() + read_csv_kwargs = _read_csv_args.ReadCsvArgs( + field_separator=_CSV_FIELD_SEPARATOR.value, + column_names=_CSV_COLUMN_NAMES.value, + column_count=_CSV_COLUMN_COUNT.value, + ).to_read_csv_kwargs() df = pd.read_csv(_DATASET_PATH.value, **read_csv_kwargs) attribute_domains = dpsynth.domain.from_yaml_file(_DOMAIN_PATH.value) diff --git a/dpsynth/bin/run_data_generation.py b/dpsynth/bin/run_data_generation.py index 2cca438..9cdbcb8 100644 --- a/dpsynth/bin/run_data_generation.py +++ b/dpsynth/bin/run_data_generation.py @@ -14,17 +14,12 @@ """Main program to launch synthetic data generation Beam jobs.""" -from collections.abc import Mapping -from typing import Any - from absl import app from absl import flags import apache_beam as beam from dpsynth import data_generation from dpsynth import domain -from dpsynth.bin import _proto_class_flag from dpsynth.dataset_descriptors import csv_descriptor -from dpsynth.dataset_descriptors import proto_descriptors from dpsynth.dataset_descriptors import tfrecord_descriptor from dpsynth.pipeline_transformations import aim from dpsynth.pipeline_transformations import input_output @@ -165,7 +160,7 @@ def get_config() -> data_generation.DataGenerationConfig: Raises: NotImplementedError: If the data format is not supported. """ - if _DATA_FORMAT.value == types.DataFormat.TFRECORD: + if _DATA_FORMAT.value == types.DataFormat.TFRECORD: descriptor = tfrecord_descriptor.get_dataset_descriptor_for_tfrecord( tfrecord_descriptor.read_tfrecords_sample(_DATASET_PATH.value), attributes=_ATTRIBUTES.value, diff --git a/dpsynth/bin/run_tabular_eval.py b/dpsynth/bin/run_tabular_eval.py index 9068858..53e2886 100644 --- a/dpsynth/bin/run_tabular_eval.py +++ b/dpsynth/bin/run_tabular_eval.py @@ -21,12 +21,9 @@ from absl import app from absl import flags import apache_beam as beam -from dpsynth.dataset_descriptors import dataset_descriptor -from dpsynth.dataset_descriptors import proto_descriptors from dpsynth.eval import tabular_eval from dpsynth.eval import types from dpsynth.pipeline_transformations import diagnostic_info -from google.protobuf import text_format import pandas as pd import pipeline_dp @@ -112,7 +109,7 @@ def _read_csv_data(): config = diagnostic_info.TabularEvalConfig( attributes=attributes, - attribute_types=[t.to_proto() for t in attribute_types], + attribute_types=[t.value for t in attribute_types], ) return original_data, synthetic_data, config @@ -133,7 +130,7 @@ def to_tuple(proto): def local_main(): """Main function for local (in-process) execution.""" assert ( - _DATA_FORMAT_STR.value == "CSV" + _DATA_FORMAT_STR.value == "csv" ), "Unsupported data format for local execution." original_data, synthetic_data, config = _read_csv_data() @@ -143,7 +140,7 @@ def local_main(): eval_report = list(eval_report_collection)[0] with open_file(_EVAL_REPORT_PATH.value, "wt") as f: - f.write(text_format.MessageToString(eval_report)) + f.write(str(eval_report)) def beam_main(): diff --git a/tests/bin/bin_entrypoints_test.py b/tests/bin/bin_entrypoints_test.py new file mode 100644 index 0000000..6c2cc87 --- /dev/null +++ b/tests/bin/bin_entrypoints_test.py @@ -0,0 +1,90 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Smoke tests for public binary entrypoints.""" + +import pathlib +import subprocess +import sys +import tempfile + +from absl.testing import absltest +from absl.testing import parameterized + + +_BINARIES = ( + 'comparison.py', + 'main.py', + 'run_data_generation.py', + 'run_data_generation_from_model.py', + 'run_tabular_eval.py', +) + + +class BinEntrypointsTest(parameterized.TestCase): + + @parameterized.parameters(*_BINARIES) + def test_help_succeeds(self, binary_name: str): + repo_root = pathlib.Path(__file__).parents[2] + result = subprocess.run( + [ + sys.executable, + str(repo_root / 'dpsynth' / 'bin' / binary_name), + '--help', + ], + cwd=repo_root, + capture_output=True, + text=True, + check=False, + ) + output = result.stdout + result.stderr + self.assertNotIn('Traceback', output) + self.assertIn('flags:', output) + + def test_run_tabular_eval_local_csv(self): + repo_root = pathlib.Path(__file__).parents[2] + with tempfile.TemporaryDirectory() as tmpdir: + tmp_path = pathlib.Path(tmpdir) + original_data = tmp_path / 'original.csv' + synthetic_data = tmp_path / 'synthetic.csv' + eval_report = tmp_path / 'eval_report.pb' + original_data.write_text('cat\nA\nB\n', encoding='utf-8') + synthetic_data.write_text('cat\nA\nC\n', encoding='utf-8') + + result = subprocess.run( + [ + sys.executable, + str(repo_root / 'dpsynth' / 'bin' / 'run_tabular_eval.py'), + f'--original_data_path={original_data}', + f'--synthetic_data_path={synthetic_data}', + f'--eval_report_path={eval_report}', + '--data_format=csv', + '--use_beam=false', + ], + cwd=repo_root, + capture_output=True, + text=True, + check=False, + ) + self.assertEqual( + result.returncode, + 0, + msg=f'run_tabular_eval.py failed:\n{result.stderr}', + ) + self.assertTrue(eval_report.exists()) + self.assertNotEmpty(eval_report.read_text(encoding='utf-8')) + + +if __name__ == '__main__': + absltest.main()