mostlyai-engine 2.3.3__tar.gz → 2.5.0__tar.gz

mostlyai_engine-2.5.0/LICENSE_HEADER

@@ -0,0 +1,13 @@
+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.

{mostlyai_engine-2.3.3 → mostlyai_engine-2.5.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mostlyai-engine
-Version: 2.3.3
+Version: 2.5.0
 Summary: Synthetic Data Engine
 Project-URL: homepage, https://github.com/mostly-ai/mostlyai-engine
 Project-URL: repository, https://github.com/mostly-ai/mostlyai-engine
@@ -8,6 +8,7 @@ Project-URL: documentation, https://mostly-ai.github.io/mostlyai-engine/
 Author-email: MOSTLY AI <dev@mostly.ai>
 License-Expression: Apache-2.0
 License-File: LICENSE
+License-File: LICENSE_HEADER
 Classifier: Development Status :: 5 - Production/Stable
 Classifier: Intended Audience :: Developers
 Classifier: Intended Audience :: Financial and Insurance Industry
@@ -17,13 +18,12 @@ Classifier: Intended Audience :: Science/Research
 Classifier: Intended Audience :: Telecommunications Industry
 Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Operating System :: OS Independent
-Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
 Classifier: Topic :: Software Development :: Libraries
 Classifier: Typing :: Typed
-Requires-Python: >=3.10
+Requires-Python: <3.14,>=3.11
 Requires-Dist: accelerate>=1.5.0
 Requires-Dist: datasets>=3.0.0
 Requires-Dist: huggingface-hub[hf-xet]>=0.30.2
@@ -38,15 +38,15 @@ Requires-Dist: pyarrow>=16.0.0
 Requires-Dist: scikit-learn>=1.4.0
 Requires-Dist: setuptools>=77.0.3
 Requires-Dist: tokenizers>=0.21.0
-Requires-Dist: torch<2.9.0,>=2.8.0
-Requires-Dist: torchaudio<2.9.0,>=2.8.0
-Requires-Dist: torchvision<0.24.0,>=0.23.0
-Requires-Dist: transformers>=4.55.0
+Requires-Dist: torch<2.10.0,>=2.9.0
+Requires-Dist: torchaudio<2.10.0,>=2.9.0
+Requires-Dist: torchvision<0.25.0,>=0.24.0
+Requires-Dist: transformers<5,>=4.55.0
 Requires-Dist: xgrammar>=0.1.21
 Provides-Extra: gpu
 Requires-Dist: bitsandbytes==0.42.0; (sys_platform == 'darwin') and extra == 'gpu'
 Requires-Dist: bitsandbytes>=0.45.5; (sys_platform == 'linux') and extra == 'gpu'
-Requires-Dist: vllm==0.10.2; (sys_platform == 'linux' or sys_platform == 'darwin') and extra == 'gpu'
+Requires-Dist: vllm==0.12.0; (sys_platform == 'linux' or sys_platform == 'darwin') and extra == 'gpu'
 Description-Content-Type: text/markdown
 
 # Synthetic Data Engine 💎
@@ -116,7 +116,7 @@ uv pip install -U 'mostlyai-engine[gpu]'
 On Linux, one can explicitly install the CPU-only variant of torch together with `mostlyai-engine`:
 
 ```bash
-uv pip install -U torch==2.8.0+cpu torchvision==0.23.0+cpu mostlyai-engine --extra-index-url https://download.pytorch.org/whl/cpu
+uv pip install -U torch==2.9.1+cpu torchvision==0.24.1+cpu mostlyai-engine --extra-index-url https://download.pytorch.org/whl/cpu
 ```
 
 ## TabularARGN for Flat Data

{mostlyai_engine-2.3.3 → mostlyai_engine-2.5.0}/README.md

@@ -65,7 +65,7 @@ uv pip install -U 'mostlyai-engine[gpu]'
 On Linux, one can explicitly install the CPU-only variant of torch together with `mostlyai-engine`:
 
 ```bash
-uv pip install -U torch==2.8.0+cpu torchvision==0.23.0+cpu mostlyai-engine --extra-index-url https://download.pytorch.org/whl/cpu
+uv pip install -U torch==2.9.1+cpu torchvision==0.24.1+cpu mostlyai-engine --extra-index-url https://download.pytorch.org/whl/cpu
 ```
 
 ## TabularARGN for Flat Data

{mostlyai_engine-2.3.3 → mostlyai_engine-2.5.0}/mostlyai/engine/__init__.py

@@ -34,7 +34,7 @@ __all__ = [
     "TabularARGN",
     "LanguageModel",
 ]
-__version__ = "2.3.3"
+__version__ = "2.5.0"
 
 # suppress specific warning related to os.fork() in multi-threaded processes
 warnings.filterwarnings("ignore", category=DeprecationWarning, message=".*multi-threaded.*fork.*")

{mostlyai_engine-2.3.3 → mostlyai_engine-2.5.0}/mostlyai/engine/_common.py

@@ -273,8 +273,9 @@ def safe_convert_datetime(values: pd.Series, date_only: bool = False) -> pd.Seri
         utc=True,
         dayfirst=False,  # assume 1/3/2020 is Jan 3
     )
-    # check whether firstday=True yields less non-NA, and if so, switch to using that flag
-    if values_parsed_fixed.isna().sum() > values.isna().sum():
+    has_slash_dates = values.astype("string").str.contains(r"^\s*\d{1,2}/\d{1,2}/\d{4}(?:\s|$)", regex=True).any()
+    # some mixed-format slash dates are interpreted more reliably with dayfirst=True
+    if has_slash_dates and values_parsed_fixed.isna().sum() > values.isna().sum():
         values_parsed_fixed_dayfirst = pd.to_datetime(
             values,
             errors="coerce",  # silently map invalid dates to NA
@@ -937,6 +938,7 @@ def impute_from_non_nan_distribution(values: pd.Series, column_stats: dict) -> t
     Returns:
         tuple[pd.Series, pd.Series]: The series with imputed values and the mask of NaNs.
     """
+    values = values.copy()
     nan_mask = values.isna()
     vc = values.value_counts(normalize=True)
     if vc.empty:
@@ -944,7 +946,8 @@ def impute_from_non_nan_distribution(values: pd.Series, column_stats: dict) -> t
     probs = vc.values
     categories = vc.index
     # NOTE: an alternative will be to use the largest remainder method
-    values[nan_mask] = np.random.choice(categories, size=nan_mask.sum(), p=probs)
+    if nan_mask.any():
+        values.loc[nan_mask] = np.random.choice(categories, size=nan_mask.sum(), p=probs)
     return values, nan_mask.astype(int)
 
 

{mostlyai_engine-2.3.3 → mostlyai_engine-2.5.0}/mostlyai/engine/_encoding_types/tabular/character.py

@@ -102,7 +102,9 @@ def encode_character(values: pd.Series, stats: dict, _: pd.Series | None = None)
     df_split = split_sub_columns_character(values, max_string_length)
     for idx in range(max_string_length):
         sub_col = f"P{idx}"
-        np_codes = np.array(pd.Categorical(df_split[sub_col], categories=stats["codes"][sub_col]).codes)
+        categories = list(stats["codes"][sub_col].keys())
+        values_at_pos = df_split[sub_col].where(df_split[sub_col].isin(categories), UNKNOWN_TOKEN)
+        np_codes = np.array(pd.Categorical(values_at_pos, categories=categories).codes)
         np.place(np_codes, np_codes == -1, 0)
         df_split[sub_col] = np_codes
     if stats["has_nan"]:

{mostlyai_engine-2.3.3 → mostlyai_engine-2.5.0}/mostlyai/engine/_encoding_types/tabular/datetime.py

@@ -234,9 +234,10 @@ def decode_datetime(df_encoded: pd.DataFrame, stats: dict):
     d = df_encoded["day"] + stats["min_values"]["day"]
     # fix invalid dates by setting these to last day of month
     is_leap = y.apply(lambda x: calendar.isleap(x))
-    d[is_leap & (m == 2) & (d > 29)] = 29
-    d[~is_leap & (m == 2) & (d > 28)] = 28
-    d[((m == 4) | (m == 6) | (m == 9) | (m == 11)) & (d > 30)] = 30
+    d = d.copy()
+    d.loc[is_leap & (m == 2) & (d > 29)] = 29
+    d.loc[~is_leap & (m == 2) & (d > 28)] = 28
+    d.loc[((m == 4) | (m == 6) | (m == 9) | (m == 11)) & (d > 30)] = 30
     # concatenate to datetime string
     y = y.astype(str)
     m = m.astype(str).str.zfill(2)
@@ -271,7 +272,7 @@ def decode_datetime(df_encoded: pd.DataFrame, stats: dict):
         # set all values to NaN if no valid values were present
         values[df_encoded["nan"] == 0] = pd.NA
     # convert from string to datetime
-    values = pd.to_datetime(values)
+    values = pd.to_datetime(values).astype("datetime64[ns]")
     if not stats["has_time"]:
-        values = pd.to_datetime(values.dt.date)
+        values = pd.to_datetime(values.dt.date).astype("datetime64[ns]")
     return values

{mostlyai_engine-2.3.3 → mostlyai_engine-2.5.0}/mostlyai/engine/_language/engine/vllm_engine.py

@@ -27,11 +27,10 @@ from peft import PeftConfig
 from pydantic import BaseModel
 from transformers import AutoConfig, AutoTokenizer
 from vllm import LLM, SamplingParams
-from vllm.config import _get_and_verify_max_len
 from vllm.distributed import cleanup_dist_env_and_memory
 from vllm.inputs.data import TokensPrompt
 from vllm.lora.request import LoRARequest
-from vllm.sampling_params import GuidedDecodingParams
+from vllm.sampling_params import StructuredOutputsParams
 
 from mostlyai.engine._language.common import is_bf16_supported
 from mostlyai.engine._language.engine.base import EngineMetrics, LanguageEngine
@@ -80,8 +79,11 @@ class VLLMEngine(LanguageEngine):
 
         model_path = str(model_path)
         self._lora_request = LoRARequest("adapter", 1, model_path)
-        config_max_model_len = _get_and_verify_max_len(
-            base_config, tokenizer_config=None, max_model_len=None, disable_sliding_window=False, sliding_window=None
+        # Get max model length from config (different models use different attribute names)
+        config_max_model_len = getattr(
+            base_config,
+            "max_position_embeddings",
+            getattr(base_config, "n_positions", getattr(base_config, "max_sequence_length", 2048)),
         )
 
         self.llm = LLM(
@@ -136,18 +138,18 @@ class VLLMEngine(LanguageEngine):
 
         sampling_params = []
         for i in range(actual_batch_size):
-            guided_decoding = None
+            structured_outputs = None
             if effective_schemas and i < len(effective_schemas):
-                # Convert Pydantic model to JSON schema for guided decoding
+                # Convert Pydantic model to JSON schema for structured output
                 schema_dict = effective_schemas[i].model_json_schema()
-                guided_decoding = GuidedDecodingParams(json=schema_dict)
+                structured_outputs = StructuredOutputsParams(json=schema_dict)
 
             sampling_params.append(
                 SamplingParams(
                     max_tokens=self.max_new_tokens,
                     temperature=sampling_temperature,
                     top_p=sampling_top_p,
-                    guided_decoding=guided_decoding,
+                    structured_outputs=structured_outputs,
                 )
             )
         t_generate = time.time()

{mostlyai_engine-2.3.3 → mostlyai_engine-2.5.0}/mostlyai/engine/_language/xgrammar_utils.py

@@ -75,14 +75,23 @@ def create_schemas(
     numeric_fields = field_types.get(ModelEncodingType.language_numeric, [])
     datetime_fields = field_types.get(ModelEncodingType.language_datetime, [])
     cache = {}
+
+    def _normalize_seed_value(seed_value):
+        return None if pd.isna(seed_value) else seed_value
+
     for _, seed_row in seed_df.iterrows():
-        cache_key = hash(tuple(sorted([(field_name, str(seed_value)) for field_name, seed_value in seed_row.items()])))
+        normalized_seed_items = [
+            (field_name, _normalize_seed_value(seed_value)) for field_name, seed_value in seed_row.items()
+        ]
+        cache_key = hash(
+            tuple(sorted([(field_name, str(seed_value)) for field_name, seed_value in normalized_seed_items]))
+        )
         if cache_key in cache:
             yield cache[cache_key]
             continue
         model_dict = {}
         if not seed_row.empty:
-            model_dict |= {field_name: (Literal[seed_value], ...) for field_name, seed_value in seed_row.items()}  # type: ignore[valid-type]
+            model_dict |= {field_name: (Literal[seed_value], ...) for field_name, seed_value in normalized_seed_items}  # type: ignore[valid-type]
         for field_name in unseeded_fields:
             if field_name in categorical_fields:
                 categories = stats["columns"][field_name]["categories"]

{mostlyai_engine-2.3.3 → mostlyai_engine-2.5.0}/mostlyai/engine/_tabular/argn.py

@@ -164,10 +164,11 @@ def _regressor_heuristic(id: str, model_size: ModelSizeOrUnits, dim_input: int,
 def _flat_context_heuristic(id: str, model_size: ModelSizeOrUnits, dim_input: int) -> list[int]:
     if isinstance(model_size, dict):
         return model_size[id]
-    model_size_layers = dict(S=[8], M=[64], L=[128])
+    model_size_layers = dict(S=[2], M=[8], L=[32])
     layers = model_size_layers[model_size]
     coefficient = round(np.log(max(dim_input, np.e)))
     dims = [unit * coefficient for unit in layers]
+    _LOG.info(f"[ARGN] flat context heuristic: {dim_input=} -> {dims}")
     return dims
 
 
@@ -176,10 +177,11 @@ def _sequential_context_heuristic(
 ) -> list[int]:
     if isinstance(model_size, dict):
         return model_size[id]
-    model_size_layers = dict(S=[8], M=[32], L=[64, 64])
+    model_size_layers = dict(S=[4], M=[16], L=[64, 64])
     layers = model_size_layers[model_size]
     coefficient = round(np.log(max(dim_input * seq_len_median, np.e)))
     dims = [unit * coefficient for unit in layers]
+    _LOG.info(f"[ARGN] sequential context heuristic: {dim_input=} x {seq_len_median=} -> {dims}")
     return dims
 
 
@@ -190,6 +192,7 @@ def _history_heuristic(id: str, model_size: ModelSizeOrUnits, dim_input: int, se
     layers = model_size_layers[model_size]
     coefficient = round(np.log(max(dim_input * seq_len_median, np.e)))
     dims = [unit * coefficient for unit in layers]
+    _LOG.info(f"[ARGN] history heuristic: {dim_input=} x {seq_len_median=} -> {dims}")
     return dims
 
 

{mostlyai_engine-2.3.3 → mostlyai_engine-2.5.0}/mostlyai/engine/_tabular/common.py

@@ -197,7 +197,7 @@ def prepare_context_inputs(
     # Build flat context inputs (CTXFLT/*)
     ctxflt_inputs = {
         col: torch.unsqueeze(
-            torch.as_tensor(ctx_encoded[col].to_numpy(), device=device).type(torch.int),
+            torch.as_tensor(ctx_encoded[col].to_numpy(copy=True), device=device).type(torch.int),
             dim=-1,
         )
         for col in ctx_encoded.columns

{mostlyai_engine-2.3.3 → mostlyai_engine-2.5.0}/mostlyai/engine/_tabular/generation.py

@@ -640,10 +640,9 @@ def decode_buffered_samples(
             keys=keys,
             key_name=tgt_context_key,
         )
-        df_syn = df_syn.drop(
-            columns=[c for c in df_syn.columns if c.startswith(POSITIONAL_COLUMN)],
-            axis=1,
-        ).reset_index(drop=True)
+        df_syn = df_syn.drop(columns=[c for c in df_syn.columns if c.startswith(POSITIONAL_COLUMN)]).reset_index(
+            drop=True
+        )
     else:
         data, seed_data = zip(*buffer.buffer)
         df_syn = pd.concat(data, axis=0).reset_index(drop=True)
@@ -820,6 +819,7 @@ def generate(
 
         if not enable_flexible_generation:
             check_column_order(gen_column_order, trn_column_order)
+
         _LOG.info(f"{rare_category_replacement_method=}")
         rare_token_fixed_probs = fix_rare_token_probs(tgt_stats, rare_category_replacement_method)
         imputation_fixed_probs = _fix_imputation_probs(tgt_stats, imputation)
@@ -1223,7 +1223,9 @@ def generate(
                 # Use context inputs prepared earlier
                 x = ctx_inputs
                 fixed_values = {
-                    col: torch.as_tensor(seed_batch_encoded[col].to_numpy(), device=model.device).type(torch.int)
+                    col: torch.as_tensor(seed_batch_encoded[col].to_numpy(copy=True), device=model.device).type(
+                        torch.int
+                    )
                     for col in seed_batch_encoded.columns
                     if col in tgt_sub_columns
                 }

{mostlyai_engine-2.3.3 → mostlyai_engine-2.5.0}/mostlyai/engine/_tabular/probability.py

@@ -245,6 +245,7 @@ def _generate_marginal_probs(
     tgt_stats: dict,
     seed_columns: list[str],
     device: torch.device,
+    n_samples: int,
     ctx_data: pd.DataFrame | None = None,
     ctx_stats: dict | None = None,
     fixed_probs: dict | None = None,
@@ -259,6 +260,7 @@ def _generate_marginal_probs(
         tgt_stats: Target statistics
         seed_columns: Seed column names in original format, in correct order
         device: Device for computation
+        n_samples: Number of samples to generate probabilities for
         ctx_data: Optional context data
         ctx_stats: Optional context statistics (required if ctx_data provided)
         fixed_probs: Optional fixed probabilities for rare token handling
@@ -266,7 +268,6 @@ def _generate_marginal_probs(
     Returns:
         DataFrame of shape (n_samples, cardinality) with probabilities and column names
     """
-    n_samples = len(seed_encoded)
     target_stats = tgt_stats["columns"][target_column]
 
     # Build fixed_values dict from seed_encoded
@@ -362,15 +363,14 @@ def predict_proba(
         )
     )
 
-    # Get seed column names (needed for column order check and _generate_marginal_probs)
     seed_columns = list(seed_data.columns)
 
-    # Check column order when flexible generation is disabled
     if not enable_flexible_generation:
         seed_columns_argn = get_argn_column_names(tgt_stats["columns"], seed_columns)
         target_columns_argn = get_argn_column_names(tgt_stats["columns"], target_columns)
-        gen_column_order = seed_columns_argn + target_columns_argn
-        check_column_order(gen_column_order, all_columns)
+        columns_to_check = seed_columns_argn + target_columns_argn
+        expected_order = [col for col in all_columns if col in columns_to_check]
+        check_column_order(columns_to_check, expected_order)
 
     # Encode seed data (features to condition on) - common for both single and multi-target
     # seed_data should NOT include any target columns
@@ -407,6 +407,7 @@ def predict_proba(
         tgt_stats=tgt_stats,
         seed_columns=seed_columns,
         device=device,
+        n_samples=n_samples,
         ctx_data=ctx_data,
         ctx_stats=ctx_stats,
         fixed_probs=fixed_probs,
@@ -451,10 +452,10 @@ def predict_proba(
         # Build DataFrames for each combo with actual values, then concatenate
         combo_dfs = []
         for combo_idx, prev_combo in enumerate(prev_combos):
-            # Copy extended_seed for this combo
-            df = extended_seed.copy()
+            # Build data dict starting with columns from extended_seed
+            data = {col: extended_seed[col].values for col in extended_seed.columns}
 
-            # Add previous target columns with actual values (no dummy values)
+            # Add previous target columns with actual values
             for i in range(target_idx):
                 prev_target_col = target_columns[i]
                 encoded_val = prev_combo[i]
@@ -466,8 +467,10 @@ def predict_proba(
                         argn_column=prev_target_stats[ARGN_COLUMN],
                         argn_sub_column=sub_col_key,
                     )
-                    df[full_sub_col_name] = encoded_val
+                    data[full_sub_col_name] = encoded_val
 
+            # Create DataFrame with explicit row count
+            df = pd.DataFrame(data, index=range(n_samples))
             combo_dfs.append(df)
 
         # Concatenate all combo DataFrames into single batch
@@ -490,6 +493,7 @@ def predict_proba(
             tgt_stats=tgt_stats,
             seed_columns=extended_seed_columns,
             device=device,
+            n_samples=n_samples * num_prev_combos,
             ctx_data=batched_ctx_data,
             ctx_stats=ctx_stats,
             fixed_probs=fixed_probs,

{mostlyai_engine-2.3.3 → mostlyai_engine-2.5.0}/pyproject.toml

@@ -1,9 +1,9 @@
 [project]
 name = "mostlyai-engine"
-version = "2.3.3"
+version = "2.5.0"
 description = "Synthetic Data Engine"
 authors = [{ name = "MOSTLY AI", email = "dev@mostly.ai" }]
-requires-python = ">=3.10"
+requires-python = ">=3.11,<3.14"
 readme = "README.md"
 license = "Apache-2.0"
 classifiers = [
@@ -14,7 +14,6 @@ classifiers = [
     "Intended Audience :: Financial and Insurance Industry",
     "Intended Audience :: Healthcare Industry",
     "Intended Audience :: Telecommunications Industry",
-    "Programming Language :: Python :: 3.10",
     "Programming Language :: Python :: 3.11",
     "Programming Language :: Python :: 3.12",
     "Programming Language :: Python :: 3.13",
@@ -25,7 +24,7 @@ classifiers = [
 ]
 
 dependencies = [
-    "setuptools>=77.0.3",  # similar to vllm 0.10.1.1
+    "setuptools>=77.0.3",
     "numpy>=2.0.0",
     "pandas>=2.2.0",
     "pyarrow>=16.0.0",
@@ -33,24 +32,24 @@ dependencies = [
     "scikit-learn>=1.4.0",
     "psutil>=5.9.5,<6",  # upgrade when colab psutil is updated
     "tokenizers>=0.21.0",
-    "transformers>=4.55.0",  # updated (removed ,<4.54.0), similar to vllm 0.10.1.1
+    "transformers>=4.55.0,<5",  # keep <5 for vllm==0.12 compatibility
     "datasets>=3.0.0",
     "accelerate>=1.5.0",
     "peft>=0.12.0",
     "huggingface-hub[hf-xet]>=0.30.2",
     "opacus>=1.5.4",
-    "xgrammar>=0.1.21",  # for vllm 0.10.1.1 compatibility (=0.1.21)
+    "xgrammar>=0.1.21",
     "json-repair>=0.47.0",
-    "torch>=2.8.0,<2.9.0",  # for vllm 0.10.2 compatibility
-    "torchaudio>=2.8.0,<2.9.0",  # for vllm 0.10.2 compatibility
-    "torchvision>=0.23.0,<0.24.0"  # for vllm 0.10.2 compatibility
+    "torch>=2.9.0,<2.10.0",
+    "torchaudio>=2.9.0,<2.10.0",
+    "torchvision>=0.24.0,<0.25.0"
 ]
 
 [project.optional-dependencies]
 gpu = [
     "bitsandbytes==0.42.0; sys_platform == 'darwin'",
     "bitsandbytes>=0.45.5; sys_platform == 'linux'",
-    "vllm==0.10.2; sys_platform == 'linux' or sys_platform == 'darwin'",
+    "vllm==0.12.0; sys_platform == 'linux' or sys_platform == 'darwin'",
 ]
 
 [dependency-groups]
@@ -95,8 +94,9 @@ requires = ["hatchling", "hatch-vcs"]
 build-backend = "hatchling.build"
 
 [tool.ruff]
-target-version = "py310"
+target-version = "py311"
 line-length = 120
+extend-exclude = ["*.ipynb"]
 [tool.ruff.format]
 exclude = ["examples/*.ipynb"]
 [tool.ruff.lint]