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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mostlyai-engine
-Version: 2.3.3
+Version: 2.4.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
@@ -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: 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>=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.4.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.4.0}/mostlyai/engine/__init__.py

@@ -34,7 +34,7 @@ __all__ = [
     "TabularARGN",
     "LanguageModel",
 ]
-__version__ = "2.3.3"
+__version__ = "2.4.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.4.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.4.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.4.0}/mostlyai/engine/_tabular/generation.py

@@ -820,6 +820,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)

{mostlyai_engine-2.3.3 → mostlyai_engine-2.4.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.4.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "mostlyai-engine"
-version = "2.3.3"
+version = "2.4.0"
 description = "Synthetic Data Engine"
 authors = [{ name = "MOSTLY AI", email = "dev@mostly.ai" }]
 requires-python = ">=3.10"
@@ -25,7 +25,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 +33,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",
     "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]