chore: dedicated data pull for embeddings to fetch all events for sequential data (#192)

Author: mplatzer

mostlyai/qa/_common.py

@@ -23,6 +23,7 @@
 
 
 ACCURACY_MAX_COLUMNS = 300  # should be an even number and greater than 100
+EMBEDDINGS_MAX_COLUMNS = 300
 
 MAX_UNIVARIATE_PLOTS = 300
 MAX_BIVARIATE_TGT_PLOTS = 300

mostlyai/qa/_embeddings.py

@@ -12,22 +12,24 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import logging
 import numpy as np
 import pandas as pd
 from sklearn.decomposition import PCA
 from sklearn.preprocessing import QuantileTransformer, normalize
+from pandas.core.dtypes.common import is_numeric_dtype, is_datetime64_dtype
 
 from mostlyai.qa._common import (
-    COUNT_COLUMN,
     EMPTY_BIN,
     NA_BIN,
-    NXT_COLUMN_PREFIX,
     RARE_BIN,
-    TGT_COLUMN_PREFIX,
 )
 from mostlyai.qa.assets import load_embedder
 
 
+_LOG = logging.getLogger(__name__)
+
+
 def encode_numerics(
     syn: pd.DataFrame, trn: pd.DataFrame, hol: pd.DataFrame | None = None
 ) -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame | None]:
@@ -128,36 +130,15 @@ def encode_strings(
 
 
 def encode_data(
-    syn: pd.DataFrame, trn: pd.DataFrame, hol: pd.DataFrame | None = None
+    syn_data: pd.DataFrame, trn_data: pd.DataFrame, hol_data: pd.DataFrame | None = None
 ) -> tuple[np.ndarray, np.ndarray, np.ndarray | None]:
     """
     Encode all columns corresponding to their data type.
     """
-    tgt_cols = [c for c in trn.columns if c.startswith(TGT_COLUMN_PREFIX)]
-    nxt_cols = [c for c in trn.columns if c.startswith(NXT_COLUMN_PREFIX)]
-    cnt_col = f"{TGT_COLUMN_PREFIX}{COUNT_COLUMN}"
-    cnt_cols = [cnt_col] if cnt_col in trn.columns else []
+    _LOG.info("encode datasets for embeddings")
     # split into numeric and string columns
-    num_dat_cols = [
-        col for col in tgt_cols if pd.api.types.is_numeric_dtype(trn[col]) or pd.api.types.is_datetime64_dtype(trn[col])
-    ]
-    string_cols = [col for col in tgt_cols if col not in num_dat_cols]
-    # keep TGT data
-    syn_data = syn[tgt_cols]
-    trn_data = trn[tgt_cols]
-    hol_data = hol[tgt_cols] if hol is not None else None
-    # append NXT data with TGT data to increase data coverage
-    if len(nxt_cols) > 0:
-        syn_nxt = syn[cnt_cols + nxt_cols]
-        syn_nxt.columns = syn_nxt.columns.str.replace(NXT_COLUMN_PREFIX, TGT_COLUMN_PREFIX)
-        syn_data = pd.concat([syn_data, syn_nxt], axis=0)
-        trn_nxt = trn[cnt_cols + nxt_cols]
-        trn_nxt.columns = trn_nxt.columns.str.replace(NXT_COLUMN_PREFIX, TGT_COLUMN_PREFIX)
-        trn_data = pd.concat([trn_data, trn_nxt], axis=0)
-        if hol is not None:
-            hol_nxt = hol[cnt_cols + nxt_cols]
-            hol_nxt.columns = hol_nxt.columns.str.replace(NXT_COLUMN_PREFIX, TGT_COLUMN_PREFIX)
-            hol_data = pd.concat([hol_data, hol_nxt], axis=0)
+    num_dat_cols = [col for col in trn_data if is_numeric_dtype(trn_data[col]) or is_datetime64_dtype(trn_data[col])]
+    string_cols = [col for col in trn_data if col not in num_dat_cols]
     # encode numeric columns
     syn_num, trn_num, hol_num = encode_numerics(
         syn_data[num_dat_cols], trn_data[num_dat_cols], hol_data[num_dat_cols] if hol_data is not None else None

mostlyai/qa/_sampling.py

@@ -34,9 +34,11 @@
 import pandas as pd
 import pyarrow as pa
 
+from mostlyai.qa._embeddings import encode_data
 from mostlyai.qa._accuracy import bin_data
 from mostlyai.qa._common import (
     CTX_COLUMN_PREFIX,
+    EMBEDDINGS_MAX_COLUMNS,
     TGT_COLUMN_PREFIX,
     NXT_COLUMN_PREFIX,
     COUNT_COLUMN,
@@ -131,7 +133,7 @@ def prepare_data_for_accuracy(
     # harmonize dtypes
     df = df.apply(harmonize_dtype)
 
-    # coerce dtypes to trn_dtypes
+    # coerce dtypes to ori_dtypes
     for trn_col, trn_dtype in (ori_dtypes or {}).items():
         if is_numeric_dtype(trn_dtype):
             df[trn_col] = pd.to_numeric(df[trn_col], errors="coerce")
@@ -262,3 +264,82 @@ def is_timestamp_dtype(x: pd.Series) -> bool:
 def is_text_heuristic(x: pd.Series) -> bool:
     # if more than 5% of rows contain unique values -> consider as TEXT
     return x.dtype == "object" and x.value_counts().eq(1).reindex(x).mean() > 0.05
+
+
+def prepare_data_for_embeddings(
+    *,
+    syn_tgt_data: pd.DataFrame,
+    trn_tgt_data: pd.DataFrame,
+    hol_tgt_data: pd.DataFrame | None = None,
+    syn_ctx_data: pd.DataFrame | None = None,
+    trn_ctx_data: pd.DataFrame | None = None,
+    hol_ctx_data: pd.DataFrame | None = None,
+    ctx_primary_key: str | None = None,
+    tgt_context_key: str | None = None,
+    max_sample_size: int | None = None,
+) -> tuple[np.ndarray, np.ndarray, np.ndarray | None]:
+    # helper variables
+    key = tgt_context_key or None
+    hol = hol_tgt_data is not None
+
+    # filter target to context keys
+    if trn_ctx_data is not None:
+        rename_key = {ctx_primary_key: key}
+        syn_ctx_data = syn_ctx_data[[ctx_primary_key]].rename(columns=rename_key)
+        trn_ctx_data = trn_ctx_data[[ctx_primary_key]].rename(columns=rename_key)
+        hol_ctx_data = hol_ctx_data[[ctx_primary_key]].rename(columns=rename_key) if hol else None
+        syn_tgt_data = syn_tgt_data.merge(syn_ctx_data, on=key, how="inner")
+        trn_tgt_data = trn_tgt_data.merge(trn_ctx_data, on=key, how="inner")
+        hol_tgt_data = hol_tgt_data.merge(hol_ctx_data, on=key, how="inner") if hol else None
+
+    # enrich with count column
+    if tgt_context_key is not None:
+        syn_tgt_data.insert(0, COUNT_COLUMN, syn_tgt_data.groupby(key)[key].transform("size"))
+        trn_tgt_data.insert(0, COUNT_COLUMN, trn_tgt_data.groupby(key)[key].transform("size"))
+        hol_tgt_data.insert(0, COUNT_COLUMN, hol_tgt_data.groupby(key)[key].transform("size")) if hol else None
+
+    # cap to Q95 sequence length of original to avoid excessive samples per group distorting results
+    if tgt_context_key is not None:
+        q95_sequence_length = trn_tgt_data.groupby(key).size().quantile(0.95)
+        syn_tgt_data = syn_tgt_data.groupby(key).sample(frac=1).groupby(key).head(n=q95_sequence_length)
+        trn_tgt_data = trn_tgt_data.groupby(key).sample(frac=1).groupby(key).head(n=q95_sequence_length)
+        hol_tgt_data = (
+            hol_tgt_data.groupby(key).sample(frac=1).groupby(key).head(n=q95_sequence_length) if hol else None
+        )
+
+    # drop key from data as its not relevant for embeddings
+    if tgt_context_key is not None:
+        syn_tgt_data = syn_tgt_data.drop(columns=[key])
+        trn_tgt_data = trn_tgt_data.drop(columns=[key])
+        hol_tgt_data = hol_tgt_data.drop(columns=[key]) if hol else None
+
+    # draw equally sized samples for fair 3-way comparison
+    max_sample_size = min(
+        max_sample_size or float("inf"),
+        len(syn_tgt_data),
+        len(trn_tgt_data),
+        len(hol_tgt_data) if hol_tgt_data is not None else float("inf"),
+    )
+    syn_tgt_data = syn_tgt_data.sample(n=max_sample_size)
+    trn_tgt_data = trn_tgt_data.sample(n=max_sample_size)
+    hol_tgt_data = hol_tgt_data.sample(n=max_sample_size) if hol else None
+
+    # limit to same columns
+    trn_cols = list(trn_tgt_data.columns)[:EMBEDDINGS_MAX_COLUMNS]
+    syn_tgt_data = syn_tgt_data[trn_cols]
+    trn_tgt_data = trn_tgt_data[trn_cols]
+    hol_tgt_data = hol_tgt_data[trn_cols] if hol else None
+
+    # harmonize dtypes
+    syn_tgt_data = syn_tgt_data.apply(harmonize_dtype)
+    trn_tgt_data = trn_tgt_data.apply(harmonize_dtype)
+    hol_tgt_data = hol_tgt_data.apply(harmonize_dtype) if hol else None
+
+    # encode data
+    syn_embeds, trn_embeds, hol_embeds = encode_data(
+        syn_data=syn_tgt_data,
+        trn_data=trn_tgt_data,
+        hol_data=hol_tgt_data,
+    )
+
+    return syn_embeds, trn_embeds, hol_embeds

mostlyai/qa/reporting.py

@@ -20,7 +20,7 @@
 import pandas as pd
 from pandas.core.dtypes.common import is_numeric_dtype, is_datetime64_dtype
 
-from mostlyai.qa import _distances, _similarity, _html_report, _embeddings
+from mostlyai.qa import _distances, _similarity, _html_report
 from mostlyai.qa._accuracy import (
     bin_data,
     binning_data,
@@ -50,6 +50,7 @@
 from mostlyai.qa._sampling import (
     prepare_data_for_accuracy,
     prepare_data_for_coherence,
+    prepare_data_for_embeddings,
 )
 from mostlyai.qa._common import (
     determine_data_size,
@@ -263,17 +264,16 @@ def report(
         progress.update(completed=15, total=100)
 
         _LOG.info("calculate embeddings")
-        # ensure that embeddings are all equal size for a fair 3-way comparison
-        max_sample_size_embeddings_final = min(
-            max_sample_size_embeddings or float("inf"),
-            syn_sample_size,
-            trn_sample_size,
-            hol_sample_size or float("inf"),
-        )
-        syn_embeds, trn_embeds, hol_embeds = _embeddings.encode_data(
-            syn=syn.head(max_sample_size_embeddings_final),
-            trn=trn.head(max_sample_size_embeddings_final),
-            hol=hol.head(max_sample_size_embeddings_final) if hol is not None else None,
+        syn_embeds, trn_embeds, hol_embeds = prepare_data_for_embeddings(
+            syn_tgt_data=syn_tgt_data,
+            trn_tgt_data=trn_tgt_data,
+            hol_tgt_data=hol_tgt_data,
+            syn_ctx_data=syn_ctx_data,
+            trn_ctx_data=trn_ctx_data,
+            hol_ctx_data=hol_ctx_data,
+            ctx_primary_key=ctx_primary_key,
+            tgt_context_key=tgt_context_key,
+            max_sample_size=max_sample_size_embeddings,
         )
         progress.update(completed=20, total=100)
 

tests/end_to_end/test_report.py

@@ -304,4 +304,3 @@ def generate_dates(start_date, end_date, num_samples):
             "Expected a warning about dtype mismatch for column 'dt'"
         )
     assert statistics.accuracy.overall > 0.6
-    assert 0.2 < statistics.similarity.discriminator_auc_training_synthetic < 0.8

tests/unit/test_html_report.py

@@ -34,9 +34,9 @@ def test_generate_store_report(tmp_path, cols, workspace):
     acc_seqs_per_cat = pd.DataFrame({"column": acc_uni["column"], "accuracy": 0.5, "accuracy_max": 0.5})
     corr_trn = _accuracy.calculate_correlations(acc_trn)
     syn_embeds, trn_embeds, hol_embeds = _embeddings.encode_data(
-        syn=syn,
-        trn=trn,
-        hol=hol,
+        syn_data=syn,
+        trn_data=trn,
+        hol_data=hol,
     )
     sim_cosine_trn_hol, sim_cosine_trn_syn = _similarity.calculate_cosine_similarities(
         syn_embeds=syn_embeds,