MSD-1124: harmonize max_sample_size_embeddings + remove default max_sample_size

Author: mplatzer

README.md

@@ -67,8 +67,8 @@ def report(
     report_subtitle: str = "",
     report_credits: str = REPORT_CREDITS,
     report_extra_info: str = "",
-    max_sample_size_accuracy: int = MAX_SAMPLE_SIZE_ACCURACY,
-    max_sample_size_embeddings: int = MAX_SAMPLE_SIZE_EMBEDDINGS,
+    max_sample_size_accuracy: int | None = None,
+    max_sample_size_embeddings: int | None = None,
     statistics_path: str | Path | None = None,
     on_progress: ProgressCallback | None = None,
 ) -> tuple[Path, dict | None]:

examples/benchmark.ipynb

@@ -58,7 +58,6 @@
     "                    syn_tgt_data=syn,\n",
     "                    trn_tgt_data=tgt,\n",
     "                    hol_tgt_data=hol,\n",
-    "                    max_sample_size_embeddings=50_000,\n",
     "                )\n",
     "                row = pd.json_normalize(metrics, sep=\"_\")\n",
     "                row.insert(0, \"dataset\", dataset)\n",

src/mostlyai/qa/common.py

@@ -23,8 +23,6 @@
 _LOG = logging.getLogger(__name__)
 
 
-MAX_SAMPLE_SIZE_ACCURACY = 100_000
-MAX_SAMPLE_SIZE_EMBEDDINGS = 10_000
 ACCURACY_MAX_COLUMNS = 300  # should be an even number and greater than 100
 
 MAX_UNIVARIATE_PLOTS = 300

src/mostlyai/qa/report.py

@@ -41,8 +41,6 @@
     ProgressCallback,
     PrerequisiteNotMetError,
     check_min_sample_size,
-    MAX_SAMPLE_SIZE_ACCURACY,
-    MAX_SAMPLE_SIZE_EMBEDDINGS,
     add_tqdm,
     NXT_COLUMN,
     CTX_COLUMN_PREFIX,
@@ -69,8 +67,8 @@ def report(
     report_subtitle: str = "",
     report_credits: str = REPORT_CREDITS,
     report_extra_info: str = "",
-    max_sample_size_accuracy: int = MAX_SAMPLE_SIZE_ACCURACY,
-    max_sample_size_embeddings: int = MAX_SAMPLE_SIZE_EMBEDDINGS,
+    max_sample_size_accuracy: int | None = None,
+    max_sample_size_embeddings: int | None = None,
     statistics_path: str | Path | None = None,
     on_progress: ProgressCallback | None = None,
 ) -> tuple[Path, dict | None]:
@@ -225,6 +223,11 @@ def report(
         )
         on_progress(current=30, total=100)
 
+        # ensure that embeddings are all of equal size for a fair 3-way comparison
+        max_sample_size_embeddings = min(syn_sample_size, trn_sample_size)
+        if hol_sample_size != 0:
+            max_sample_size_embeddings = min(max_sample_size_embeddings, hol_sample_size)
+
         # calculate embeddings
         syn_embeds = calculate_embeddings(
             pull_data_for_embeddings(
@@ -237,21 +240,16 @@ def report(
         )
         _LOG.info(f"calculated embeddings for synthetic {syn_embeds.shape}")
         on_progress(current=50, total=100)
-        # ensure that `trn` and `hol` are of equal size
-        max_sample_size = min(max_sample_size_embeddings, trn_sample_size)
-        if hol_tgt_data is not None:
-            max_sample_size = min(max_sample_size_embeddings, hol_sample_size)
         trn_embeds = calculate_embeddings(
             pull_data_for_embeddings(
                 df_tgt=trn_tgt_data,
                 df_ctx=trn_ctx_data,
                 ctx_primary_key=ctx_primary_key,
                 tgt_context_key=tgt_context_key,
-                max_sample_size=max_sample_size,
+                max_sample_size=max_sample_size_embeddings,
             )
         )
         _LOG.info(f"calculated embeddings for training {trn_embeds.shape}")
-
         on_progress(current=60, total=100)
         if hol_tgt_data is not None:
             hol_embeds = calculate_embeddings(
@@ -260,7 +258,7 @@ def report(
                     df_ctx=hol_ctx_data,
                     ctx_primary_key=ctx_primary_key,
                     tgt_context_key=tgt_context_key,
-                    max_sample_size=max_sample_size,
+                    max_sample_size=max_sample_size_embeddings,
                 )
             )
             _LOG.info(f"calculated embeddings for holdout {hol_embeds.shape}")

src/mostlyai/qa/report_from_statistics.py

@@ -27,8 +27,6 @@
     PrerequisiteNotMetError,
     check_min_sample_size,
     add_tqdm,
-    MAX_SAMPLE_SIZE_ACCURACY,
-    MAX_SAMPLE_SIZE_EMBEDDINGS,
     check_statistics_prerequisite,
     determine_data_size,
     REPORT_CREDITS,
@@ -50,8 +48,8 @@ def report_from_statistics(
     report_subtitle: str = "",
     report_credits: str = REPORT_CREDITS,
     report_extra_info: str = "",
-    max_sample_size_accuracy: int = MAX_SAMPLE_SIZE_ACCURACY,
-    max_sample_size_embeddings: int = MAX_SAMPLE_SIZE_EMBEDDINGS,
+    max_sample_size_accuracy: int | None = None,
+    max_sample_size_embeddings: int | None = None,
     on_progress: ProgressCallback | None = None,
 ) -> Path:
     with TemporaryWorkspace() as workspace:

src/mostlyai/qa/sampling.py

@@ -70,19 +70,20 @@ def pull_data_for_accuracy(
 
     if df_ctx is not None:
         # explicit context
-        df_ctx = df_ctx.sample(frac=1).head(max_sample_size).reset_index(drop=True)
-        df_ctx = df_ctx.rename(columns={ctx_primary_key: tgt_context_key})
-        df_tgt = df_tgt.merge(df_ctx[tgt_context_key], on=tgt_context_key).reset_index(drop=True)
+        df_ctx = df_ctx.sample(frac=1).head(max_sample_size)
+        df_ctx = df_ctx.rename(columns={ctx_primary_key: tgt_context_key}).reset_index(drop=True)
+        df_tgt = df_tgt.merge(df_ctx[tgt_context_key], on=tgt_context_key, how="inner").reset_index(drop=True)
     elif tgt_context_key is not None:
         # implicit context
-        df_ctx = df_tgt[[tgt_context_key]].drop_duplicates().sample(frac=1).head(max_sample_size).reset_index(drop=True)
-        df_tgt = df_tgt.merge(df_ctx[tgt_context_key], on=tgt_context_key).reset_index(drop=True)
+        df_ctx = df_tgt[[tgt_context_key]].drop_duplicates()
+        df_ctx = df_ctx.sample(frac=1).head(max_sample_size).reset_index(drop=True)
+        df_tgt = df_tgt.merge(df_ctx[tgt_context_key], on=tgt_context_key, how="inner").reset_index(drop=True)
     else:
         # no context; flat table
-        df_ctx = pd.DataFrame({key: range(len(df_tgt))})
-        df_tgt = df_tgt.sample(frac=1).head(max_sample_size).reset_index(drop=True)
-        df_tgt[key] = df_ctx[key]
         tgt_context_key = key
+        df_tgt = df_tgt.sample(frac=1).head(max_sample_size).reset_index(drop=True)
+        df_tgt[key] = range(len(df_tgt))
+        df_ctx = df_tgt[[key]]
 
     # consistently use "__KEY" as key column
     df_ctx = df_ctx.rename(columns={tgt_context_key: key})
@@ -188,12 +189,13 @@ def pull_data_for_embeddings(
 
     if df_ctx is not None:
         # explicit context
-        df_ctx = df_ctx.sample(frac=1).head(max_sample_size).reset_index(drop=True)
-        df_ctx = df_ctx.rename(columns={ctx_primary_key: tgt_context_key})
+        df_ctx = df_ctx.sample(frac=1).head(max_sample_size)
+        df_ctx = df_ctx.rename(columns={ctx_primary_key: tgt_context_key}).reset_index(drop=True)
         df_tgt = df_tgt.merge(df_ctx[tgt_context_key], on=tgt_context_key, how="right").reset_index(drop=True)
     elif tgt_context_key is not None:
         # implicit context
-        df_ctx = df_tgt[[tgt_context_key]].drop_duplicates().sample(frac=1).head(max_sample_size).reset_index(drop=True)
+        df_ctx = df_tgt[[tgt_context_key]].drop_duplicates()
+        df_ctx = df_ctx.sample(frac=1).head(max_sample_size).reset_index(drop=True)
         df_tgt = df_tgt.merge(df_ctx[tgt_context_key], on=tgt_context_key, how="right").reset_index(drop=True)
     else:
         # no context; flat table
@@ -214,18 +216,18 @@ def row_to_string(row: pd.Series) -> str:
         # JSON to keep the string length for faster speed short
         return " ".join(row.values.astype(str))
 
-    def sequence_to_json(sequence: pd.DataFrame) -> str:
+    def sequence_to_string(sequence: pd.DataFrame) -> str:
         return ", ".join(sequence.apply(row_to_string, axis=1))
 
-    jsons = (
+    strings = (
         df_tgt.groupby(tgt_context_key)
-        .apply(sequence_to_json, include_groups=False)
+        .apply(sequence_to_string, include_groups=False)
         .sample(frac=1)
         .reset_index(drop=True)
     )
     time_elapsed = time.time() - t0
-    _LOG.info(f"finished pulling data for embeddings ({time_elapsed=:.2f}s, {jsons.shape=})")
-    return jsons
+    _LOG.info(f"finished pulling data for embeddings ({time_elapsed=:.2f}s, {strings.shape=})")
+    return strings
 
 
 def calculate_embeddings(texts: pd.Series | pd.DataFrame) -> np.ndarray:

tests/end_to_end/test_report.py

@@ -42,9 +42,9 @@ def test_report_flat(tmp_path):
     assert report_path.exists()
 
     accuracy = metrics["accuracy"]
-    assert 0.8 <= accuracy["overall"] <= 1.0
-    assert 0.8 <= accuracy["univariate"] <= 1.0
-    assert 0.8 <= accuracy["bivariate"] <= 1.0
+    assert 0.5 <= accuracy["overall"] <= 1.0
+    assert 0.5 <= accuracy["univariate"] <= 1.0
+    assert 0.5 <= accuracy["bivariate"] <= 1.0
     assert accuracy["coherence"] is None
     assert 0.8 <= accuracy["overall_max"] <= 1.0
     assert 0.8 <= accuracy["univariate_max"] <= 1.0