feat: improve reproducibility interface (#203)

Author: lukaszkolodziejczyk

mostlyai/qa/__init__.py

@@ -19,10 +19,11 @@
 from packaging.version import Version
 
 from mostlyai.qa.logging import init_logging
+from mostlyai.qa.random_state import set_random_state
 from mostlyai.qa.reporting import report
 from mostlyai.qa.reporting_from_statistics import report_from_statistics
 
-__all__ = ["report", "report_from_statistics", "init_logging"]
+__all__ = ["report", "report_from_statistics", "init_logging", "set_random_state"]
 __version__ = "1.9.5"
 
 warnings.filterwarnings("ignore", category=FutureWarning, module="phik")

mostlyai/qa/_common.py

@@ -13,8 +13,6 @@
 # limitations under the License.
 
 import logging
-import os
-import struct
 from typing import Protocol
 
 import pandas as pd
@@ -124,21 +122,3 @@ def determine_data_size(
         return len(tgt_keys)
     else:
         return len(tgt_data)
-
-
-def set_random_state(random_state: int | None = None):
-    def get_random_int_from_os() -> int:
-        # 32-bit, cryptographically secure random int from os
-        return int(struct.unpack("I", os.urandom(4))[0])
-
-    if random_state is not None:
-        _LOG.info(f"Global random_state set to `{random_state}`")
-
-    if random_state is None:
-        random_state = get_random_int_from_os()
-
-    import random
-    import numpy as np
-
-    random.seed(random_state)
-    np.random.seed(random_state)

mostlyai/qa/random_state.py

@@ -0,0 +1,36 @@
+# Copyright 2025 MOSTLY AI
+#
+# 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.
+
+import logging
+import os
+import struct
+import random
+import numpy as np
+
+_LOG = logging.getLogger(__name__)
+
+
+def set_random_state(random_state: int | None = None):
+    def get_random_int_from_os() -> int:
+        # 32-bit, cryptographically secure random int from os
+        return int(struct.unpack("I", os.urandom(4))[0])
+
+    if random_state is not None:
+        _LOG.info(f"Global random_state set to `{random_state}`")
+
+    if random_state is None:
+        random_state = get_random_int_from_os()
+
+    random.seed(random_state)
+    np.random.seed(random_state)

mostlyai/qa/reporting.py

@@ -62,7 +62,6 @@
     TGT_COLUMN_PREFIX,
     REPORT_CREDITS,
     ProgressCallbackWrapper,
-    set_random_state,
 )
 from mostlyai.qa._filesystem import Statistics, TemporaryWorkspace
 
@@ -88,7 +87,6 @@ def report(
     max_sample_size_embeddings: int | None = None,
     statistics_path: str | Path | None = None,
     update_progress: ProgressCallback | None = None,
-    random_state: int | None = None,
 ) -> tuple[Path, ModelMetrics | None]:
     """
     Generate an HTML report and metrics for assessing synthetic data quality.
@@ -123,15 +121,12 @@ def report(
         max_sample_size_embeddings: The maximum sample size for embedding calculations.
         statistics_path: The path of where to store the statistics to be used by `report_from_statistics`
         update_progress: The progress callback.
-        random_state: Seed for the random number generators.
 
     Returns:
         The path to the generated HTML report.
         Metrics instance with accuracy, similarity, and distances metrics.
     """
 
-    set_random_state(random_state)
-
     if syn_ctx_data is not None:
         if ctx_primary_key is None:
             raise ValueError("If syn_ctx_data is provided, then ctx_primary_key must also be provided.")

mostlyai/qa/reporting_from_statistics.py

@@ -33,7 +33,6 @@
     determine_data_size,
     REPORT_CREDITS,
     ProgressCallbackWrapper,
-    set_random_state,
 )
 from mostlyai.qa._filesystem import Statistics, TemporaryWorkspace
 
@@ -54,7 +53,6 @@ def report_from_statistics(
     max_sample_size_accuracy: int | None = None,
     max_sample_size_coherence: int | None = None,
     update_progress: ProgressCallback | None = None,
-    random_state: int | None = None,
 ) -> Path:
     """
     Generate an HTML report based on previously generated statistics and newly provided synthetic data samples.
@@ -72,14 +70,11 @@ def report_from_statistics(
         max_sample_size_accuracy: The maximum sample size for accuracy calculations.
         max_sample_size_coherence: The maximum sample size for coherence calculations.
         update_progress: The progress callback.
-        random_state: Seed for the random number generators.
 
     Returns:
         The path to the generated HTML report.
     """
 
-    set_random_state(random_state)
-
     with (
         TemporaryWorkspace() as workspace,
         ProgressCallbackWrapper(update_progress) as progress,

tests/end_to_end/test_report.py

@@ -304,3 +304,23 @@ def generate_dates(start_date, end_date, num_samples):
             "Expected a warning about dtype mismatch for column 'dt'"
         )
     assert statistics.accuracy.overall > 0.6
+
+
+def test_reproducibility(tmp_path):
+    statistics_path = tmp_path / "statistics"
+    syn_tgt_data = mock_data(150)
+    trn_tgt_data = mock_data(150)
+    hol_tgt_data = mock_data(150)
+    kwargs = {
+        "syn_tgt_data": syn_tgt_data,
+        "trn_tgt_data": trn_tgt_data,
+        "hol_tgt_data": hol_tgt_data,
+        "statistics_path": statistics_path,
+        "max_sample_size_accuracy": 120,
+        "max_sample_size_embeddings": 80,
+    }
+    qa.set_random_state(45)
+    _, m1 = qa.report(**kwargs)
+    qa.set_random_state(45)
+    _, m2 = qa.report(**kwargs)
+    assert m1 == m2