chore: sync 5737 from nmp

src/nemo_safe_synthesizer/artifacts/analyzers/data_checks/__init__.py

@@ -2,25 +2,25 @@
 # SPDX-License-Identifier: Apache-2.0
 
 # Make all checks available to import from `data_checks` module.
-from nemo_safe_synthesizer.artifacts.analyzers.data_checks.base import (
+from .base import (
     DataCheck,
 )
-from nemo_safe_synthesizer.artifacts.analyzers.data_checks.dataset_size import (
+from .dataset_size import (
     DatasetSizeCheck,
 )
-from nemo_safe_synthesizer.artifacts.analyzers.data_checks.high_float_precision import (
+from .high_float_precision import (
     HighFloatPrecisionCheck,
 )
-from nemo_safe_synthesizer.artifacts.analyzers.data_checks.missing_data import (
+from .missing_data import (
     MissingDataCheck,
 )
-from nemo_safe_synthesizer.artifacts.analyzers.data_checks.sparse_data import (
+from .sparse_data import (
     SparseDataCheck,
 )
-from nemo_safe_synthesizer.artifacts.analyzers.data_checks.surrounding_whitespaces import (
+from .surrounding_whitespaces import (
     SurroundingWhitespacesCheck,
 )
-from nemo_safe_synthesizer.artifacts.analyzers.data_checks.text_data import (
+from .text_data import (
     TextDataCheck,
 )
 

src/nemo_safe_synthesizer/artifacts/analyzers/data_checks/base.py

@@ -6,14 +6,14 @@
 from abc import ABC, abstractmethod
 from typing import Any
 
-from nemo_safe_synthesizer.artifacts.base.analyzer import (
+from ...base.analyzer import (
     AnalyzerContext,
     ArtifactAnalyzer,
 )
-from nemo_safe_synthesizer.artifacts.base.data_checks import (
+from ...base.data_checks import (
     DataCheckResult,
 )
-from nemo_safe_synthesizer.artifacts.base.metrics import timed
+from ...base.metrics import timed
 
 
 class DataChecksAnalyzer(ArtifactAnalyzer):

src/nemo_safe_synthesizer/artifacts/analyzers/data_checks/dataset_size.py

@@ -5,14 +5,14 @@
 
 from typing import TYPE_CHECKING
 
-from nemo_safe_synthesizer.artifacts.analyzers.data_checks.base import DataCheck
-from nemo_safe_synthesizer.artifacts.base.data_checks import (
+from ...base.data_checks import (
     DataCheckResult,
     DataCheckWarning,
 )
+from .base import DataCheck
 
 if TYPE_CHECKING:
-    from nemo_safe_synthesizer.artifacts.base.analyzer import AnalyzerContext
+    from ...base.analyzer import AnalyzerContext
 
 _RECORD_COUNT_LIMIT = 50_000
 _FIELD_COUNT_LIMIT = 30

src/nemo_safe_synthesizer/artifacts/analyzers/data_checks/high_float_precision.py

@@ -5,20 +5,22 @@
 
 from typing import TYPE_CHECKING
 
-from nemo_safe_synthesizer.artifacts.analyzers.data_checks.base import (
+from ...base.data_checks import (
+    DataCheckResult,
+    DataCheckWarning,
+)
+from ...base.fields import (
+    FieldType,
+)
+from .base import (
     DataCheck,
     plural_verb,
     warning_explain_prefix,
 )
-from nemo_safe_synthesizer.artifacts.base.data_checks import (
-    DataCheckResult,
-    DataCheckWarning,
-)
-from nemo_safe_synthesizer.artifacts.base.fields import FieldType
 
 if TYPE_CHECKING:
-    from nemo_safe_synthesizer.artifacts.base.analyzer import AnalyzerContext
-    from nemo_safe_synthesizer.artifacts.base.fields import FieldFeatures
+    from ...base.analyzer import AnalyzerContext
+    from ...base.fields import FieldFeatures
 
 _FLOAT_PRECISION_LIMIT = 6
 

src/nemo_safe_synthesizer/artifacts/analyzers/data_checks/missing_data.py

@@ -5,19 +5,19 @@
 
 from typing import TYPE_CHECKING
 
-from nemo_safe_synthesizer.artifacts.analyzers.data_checks.base import (
+from ...base.data_checks import (
+    DataCheckResult,
+    DataCheckWarning,
+)
+from .base import (
     DataCheck,
     plural_verb,
     warning_explain_prefix,
 )
-from nemo_safe_synthesizer.artifacts.base.data_checks import (
-    DataCheckResult,
-    DataCheckWarning,
-)
 
 if TYPE_CHECKING:
-    from nemo_safe_synthesizer.artifacts.base.analyzer import AnalyzerContext
-    from nemo_safe_synthesizer.artifacts.base.fields import FieldFeatures
+    from ...base.analyzer import AnalyzerContext
+    from ...base.fields import FieldFeatures
 
 _MISSING_PERCENTAGE_LIMIT = 50
 

src/nemo_safe_synthesizer/artifacts/analyzers/data_checks/sparse_data.py

@@ -7,24 +7,24 @@
 
 import numpy as np
 
-from nemo_safe_synthesizer.artifacts.analyzers.data_checks.base import (
-    DataCheck,
-    plural_verb,
-    warning_explain_prefix,
-)
-from nemo_safe_synthesizer.artifacts.base.data_checks import (
+from ...base.data_checks import (
     DataCheckResult,
     DataCheckWarning,
 )
-from nemo_safe_synthesizer.artifacts.base.fields import (
+from ...base.fields import (
     FieldType,
 )
+from .base import (
+    DataCheck,
+    plural_verb,
+    warning_explain_prefix,
+)
 
 if TYPE_CHECKING:
-    from nemo_safe_synthesizer.artifacts.base.analyzer import (
+    from ...base.analyzer import (
         AnalyzerContext,
     )
-    from nemo_safe_synthesizer.artifacts.base.fields import (
+    from ...base.fields import (
         FieldFeatures,
     )
 

src/nemo_safe_synthesizer/artifacts/analyzers/data_checks/surrounding_whitespaces.py

@@ -8,18 +8,18 @@
 from pandas import Series
 from pandas.core.dtypes.common import is_string_or_object_np_dtype
 
-from nemo_safe_synthesizer.artifacts.analyzers.data_checks.base import (
+from ...base.data_checks import (
+    DataCheckResult,
+    DataCheckWarning,
+)
+from .base import (
     DataCheck,
     plural_verb,
     warning_explain_prefix,
 )
-from nemo_safe_synthesizer.artifacts.base.data_checks import (
-    DataCheckResult,
-    DataCheckWarning,
-)
 
 if TYPE_CHECKING:
-    from nemo_safe_synthesizer.artifacts.base.analyzer import AnalyzerContext
+    from ...base.analyzer import AnalyzerContext
 
 _VALUES_WITH_WHITESPACE_PERCENTAGE_LIMIT = 20
 

src/nemo_safe_synthesizer/artifacts/analyzers/data_checks/text_data.py

@@ -5,19 +5,19 @@
 
 from typing import TYPE_CHECKING
 
-from nemo_safe_synthesizer.artifacts.analyzers.data_checks.base import (
+from ...base.data_checks import (
+    DataCheckResult,
+    DataCheckWarning,
+)
+from ...base.fields import FieldType
+from .base import (
     DataCheck,
     plural_verb,
     warning_explain_prefix,
 )
-from nemo_safe_synthesizer.artifacts.base.data_checks import (
-    DataCheckResult,
-    DataCheckWarning,
-)
-from nemo_safe_synthesizer.artifacts.base.fields import FieldType
 
 if TYPE_CHECKING:
-    from nemo_safe_synthesizer.artifacts.base.analyzer import AnalyzerContext
+    from ...base.analyzer import AnalyzerContext
 
 _TEXT_FIELDS_COUNT_LIMIT = 2
 

src/nemo_safe_synthesizer/artifacts/analyzers/field_features.py

@@ -12,19 +12,19 @@
 from pandas import Series
 from pandas.core.dtypes.common import is_float_dtype, is_numeric_dtype
 
-from nemo_safe_synthesizer.artifacts.base.analyzer import (
+from ..base.analyzer import (
     AnalyzerContext,
     ArtifactAnalyzer,
 )
-from nemo_safe_synthesizer.artifacts.base.fields import (
+from ..base.fields import (
     FieldFeatures,
     FieldFeaturesInfo,
     FieldType,
 )
-from nemo_safe_synthesizer.artifacts.base.metrics import timed
+from ..base.metrics import timed
 
 if TYPE_CHECKING:
-    from nemo_safe_synthesizer.artifacts.base.metadata import (
+    from ..base.metadata import (
         DatasetMetadata,
     )
 

src/nemo_safe_synthesizer/artifacts/base/analyzer.py

@@ -9,20 +9,20 @@
 from pandas import DataFrame
 from pydantic import BaseModel, ConfigDict, Field
 
-from nemo_safe_synthesizer.artifacts.base.name_anonymizer import NoopNameAnonymizer
+from .name_anonymizer import NoopNameAnonymizer
 
 if TYPE_CHECKING:
-    from nemo_safe_synthesizer.artifacts.base.data_checks import (
+    from .data_checks import (
         DataCheckResult,
     )
-    from nemo_safe_synthesizer.artifacts.base.fields import (
+    from .fields import (
         FieldFeatures,
         FieldFeaturesInfo,
     )
-    from nemo_safe_synthesizer.artifacts.base.metadata import (
+    from .metadata import (
         DatasetMetadata,
     )
-    from nemo_safe_synthesizer.artifacts.base.name_anonymizer import NameAnonymizer
+    from .name_anonymizer import NameAnonymizer
 
 
 class AnalyzerContext(BaseModel):

src/nemo_safe_synthesizer/artifacts/base/metadata.py

@@ -7,7 +7,7 @@
 
 from pydantic import Field
 
-from nemo_safe_synthesizer.config.base import NSSBaseModel
+from ...config.base import NSSBaseModel
 
 
 class FieldAttribute(StrEnum):

src/nemo_safe_synthesizer/artifacts/data_checks_processor.py

@@ -5,25 +5,15 @@
 
 from pandas import DataFrame
 
-from nemo_safe_synthesizer.artifacts.analyzers.data_checks import (
+from ..observability import get_logger
+from .analyzers.data_checks import (
     create_all_checks,
 )
-from nemo_safe_synthesizer.artifacts.analyzers.data_checks.base import (
-    DataChecksAnalyzer,
-)
-from nemo_safe_synthesizer.artifacts.analyzers.field_features import (
-    FieldFeaturesAnalyzer,
-)
-from nemo_safe_synthesizer.artifacts.base.analyzer import (
-    AnalyzerContext,
-    ArtifactAnalyzer,
-    DataFrameBasicAnalyzer,
-)
-from nemo_safe_synthesizer.artifacts.base.data_checks import (
-    DataCheckResults,
-)
-from nemo_safe_synthesizer.artifacts.base.name_anonymizer import NoopNameAnonymizer
-from nemo_safe_synthesizer.observability import get_logger
+from .analyzers.data_checks.base import DataChecksAnalyzer
+from .analyzers.field_features import FieldFeaturesAnalyzer
+from .base.analyzer import AnalyzerContext, ArtifactAnalyzer, DataFrameBasicAnalyzer
+from .base.data_checks import DataCheckResults
+from .base.name_anonymizer import NoopNameAnonymizer
 
 logger = get_logger(__name__)
 

src/nemo_safe_synthesizer/cli/artifacts.py

@@ -8,7 +8,7 @@
 
 import click
 
-from nemo_safe_synthesizer.cli.artifact_structure import BoundDir, PathT, Workdir
+from .artifact_structure import BoundDir, PathT, Workdir
 
 
 @click.group(invoke_without_command=True)

src/nemo_safe_synthesizer/cli/run.py

@@ -30,7 +30,7 @@
 )
 
 if TYPE_CHECKING:
-    from nemo_safe_synthesizer.sdk.library_builder import SafeSynthesizer
+    from ..sdk.library_builder import SafeSynthesizer
 
 
 def common_run_options(f):
@@ -216,7 +216,7 @@ def run(
     run_logger.debug("running with: ", extra={"config": config.model_dump()})
 
     with traced_user("SafeSynthesizer"):
-        from nemo_safe_synthesizer.sdk.library_builder import SafeSynthesizer
+        from ..sdk.library_builder import SafeSynthesizer
 
         ss: SafeSynthesizer = SafeSynthesizer(config=config, workdir=workdir).with_data_source(df)
         ss.run()
@@ -270,7 +270,7 @@ def run_train(
         settings=settings,
         phase="train",
     )
-    from nemo_safe_synthesizer.sdk.library_builder import SafeSynthesizer
+    from ..sdk.library_builder import SafeSynthesizer
 
     with traced_user("SafeSynthesizer"):
         SafeSynthesizer(config, workdir=workdir).with_data_source(df).process_data().train()
@@ -347,7 +347,7 @@ def run_generate(
         auto_discover_adapter=auto_discover_adapter,
         wandb_resume_job_id=wandb_resume_job_id,
     )
-    from nemo_safe_synthesizer.sdk.library_builder import SafeSynthesizer
+    from ..sdk.library_builder import SafeSynthesizer
 
     final_output_file = settings.output_file or workdir.output_file
     with traced_user("SafeSynthesizer"):

src/nemo_safe_synthesizer/cli/utils.py

@@ -272,7 +272,7 @@ def common_setup(
     config = merge_overrides(settings.config_path, synthesis_overrides)
 
     # 6. Initialize wandb (uses workdir for run ID files)
-    initialize_wandb_run(workdir, resume_job_id=wandb_resume_job_id)
+    initialize_wandb_run(workdir, resume_job_id=wandb_resume_job_id, cfg=config)
 
     return run_logger, config, df, workdir