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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mostlyai-engine
-Version: 2.4.0
+Version: 2.6.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
@@ -17,13 +17,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
@@ -32,21 +31,21 @@ Requires-Dist: json-repair>=0.47.0
 Requires-Dist: numpy>=2.0.0
 Requires-Dist: opacus>=1.5.4
 Requires-Dist: pandas>=2.2.0
-Requires-Dist: peft>=0.12.0
+Requires-Dist: peft>=0.18.2
 Requires-Dist: psutil<6,>=5.9.5
 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.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
+Requires-Dist: setuptools<81.0.0,>=77.0.3
+Requires-Dist: tokenizers>=0.21.1
+Requires-Dist: torch<2.12.0,>=2.11.0
+Requires-Dist: torchaudio<2.12.0,>=2.11.0
+Requires-Dist: torchvision<0.27.0,>=0.26.0
+Requires-Dist: transformers>=5.5.1
+Requires-Dist: xgrammar<1.0.0,>=0.1.32
 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.12.0; (sys_platform == 'linux' or sys_platform == 'darwin') and extra == 'gpu'
+Requires-Dist: vllm==0.20.0; (sys_platform == 'linux' or sys_platform == 'darwin') and extra == 'gpu'
 Description-Content-Type: text/markdown
 
 # Synthetic Data Engine 💎
@@ -113,10 +112,13 @@ or alternatively for a GPU setup (needed for LLM finetuning and inference):
 uv pip install -U 'mostlyai-engine[gpu]'
 ```
 
-On Linux, one can explicitly install the CPU-only variant of torch together with `mostlyai-engine`:
+On Linux, one can explicitly install the CPU-only variant of PyTorch together with `mostlyai-engine`:
 
 ```bash
-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
+uv pip install --index-strategy unsafe-first-match -U \
+  torch==2.11.0+cpu torchvision==0.26.0+cpu torchaudio==2.11.0+cpu \
+  mostlyai-engine \
+  --extra-index-url https://download.pytorch.org/whl/cpu
 ```
 
 ## TabularARGN for Flat Data
@@ -191,10 +193,15 @@ from sklearn.metrics import accuracy_score, roc_auc_score
 
 # predict class labels for a categorical
 predictions = argn.predict(data_test, target="income", n_draws=100, agg_fn="mode")
+# model-conditional class probabilities (same inputs as predict; target column dropped from seed)
+probabilities = argn.predict_proba(data_test, target="income")
 
 # evaluate performance
-accuracy = accuracy_score(data_test["income"], predictions)
-auc = roc_auc_score(data_test["income"], probabilities[:, 1])
+accuracy = accuracy_score(data_test["income"], predictions["income"])
+# AUC: sklearn needs binary 0/1 targets and scores for the "positive" class (here: second category)
+pos_label = probabilities.columns[1]
+y_true_bin = (data_test["income"] == pos_label).astype(int)
+auc = roc_auc_score(y_true_bin, probabilities[pos_label])
 print(f"Accuracy: {accuracy:.3f}, AUC: {auc:.3f}")
 ```
 
@@ -305,7 +312,7 @@ argn.sample(ctx_data=ctx_data)
 
 The `LanguageModel` class provides a scikit-learn-compatible interface for working with semi-structured textual data. It leverages pre-trained language models or trains lightweight LSTM models from scratch to generate synthetic text data.
 
-**Note**: The default model is `MOSTLY_AI/LSTMFromScratch-3m`, a lightweight LSTM model trained from scratch (GPU strongly recommended). You can also use pre-trained HuggingFace models by setting model to e.g. `microsoft/phi-1.5` (GPU required).
+**Note**: The default model is `MOSTLY_AI/LSTMFromScratch-3m`, a lightweight LSTM model trained from scratch (GPU strongly recommended). You can also use pretrained Hugging Face models (`model="<hub/repo>"`; GPU required). Verified checkpoints include `HuggingFaceTB/SmolLM2-135M`, `HuggingFaceTB/SmolLM3-3B`, `Qwen/Qwen3-0.6B`, and `microsoft/phi-4`.
 
 ### Model Training
 
@@ -360,5 +367,5 @@ lm.sample(
 
 Example notebooks demonstrating various use cases are available in the `examples` directory:
 - TabularARGN for flat tabular data [![Run on Colab](https://img.shields.io/badge/Open%20in-Colab-blue?logo=google-colab)](https://colab.research.google.com/github/mostly-ai/mostlyai-engine/blob/main/examples/flat.ipynb)
-- TabularARGN for sequential data [![Run on Colab](https://img.shields.io/badge/Open%20in-Colab-blue?logo=google-colab)](https://colab.research.google.com/github/mostly-ai/mostlyai-engine/blob/main/examples/language.ipynb)
-- LanguageModel for textual data [![Run on Colab](https://img.shields.io/badge/Open%20in-Colab-blue?logo=google-colab)](https://colab.research.google.com/github/mostly-ai/mostlyai-engine/blob/main/examples/sequential.ipynb)
+- TabularARGN for sequential data [![Run on Colab](https://img.shields.io/badge/Open%20in-Colab-blue?logo=google-colab)](https://colab.research.google.com/github/mostly-ai/mostlyai-engine/blob/main/examples/sequential.ipynb)
+- LanguageModel for textual data [![Run on Colab](https://img.shields.io/badge/Open%20in-Colab-blue?logo=google-colab)](https://colab.research.google.com/github/mostly-ai/mostlyai-engine/blob/main/examples/language.ipynb)

{mostlyai_engine-2.4.0 → mostlyai_engine-2.6.0}/README.md

@@ -62,10 +62,13 @@ or alternatively for a GPU setup (needed for LLM finetuning and inference):
 uv pip install -U 'mostlyai-engine[gpu]'
 ```
 
-On Linux, one can explicitly install the CPU-only variant of torch together with `mostlyai-engine`:
+On Linux, one can explicitly install the CPU-only variant of PyTorch together with `mostlyai-engine`:
 
 ```bash
-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
+uv pip install --index-strategy unsafe-first-match -U \
+  torch==2.11.0+cpu torchvision==0.26.0+cpu torchaudio==2.11.0+cpu \
+  mostlyai-engine \
+  --extra-index-url https://download.pytorch.org/whl/cpu
 ```
 
 ## TabularARGN for Flat Data
@@ -140,10 +143,15 @@ from sklearn.metrics import accuracy_score, roc_auc_score
 
 # predict class labels for a categorical
 predictions = argn.predict(data_test, target="income", n_draws=100, agg_fn="mode")
+# model-conditional class probabilities (same inputs as predict; target column dropped from seed)
+probabilities = argn.predict_proba(data_test, target="income")
 
 # evaluate performance
-accuracy = accuracy_score(data_test["income"], predictions)
-auc = roc_auc_score(data_test["income"], probabilities[:, 1])
+accuracy = accuracy_score(data_test["income"], predictions["income"])
+# AUC: sklearn needs binary 0/1 targets and scores for the "positive" class (here: second category)
+pos_label = probabilities.columns[1]
+y_true_bin = (data_test["income"] == pos_label).astype(int)
+auc = roc_auc_score(y_true_bin, probabilities[pos_label])
 print(f"Accuracy: {accuracy:.3f}, AUC: {auc:.3f}")
 ```
 
@@ -254,7 +262,7 @@ argn.sample(ctx_data=ctx_data)
 
 The `LanguageModel` class provides a scikit-learn-compatible interface for working with semi-structured textual data. It leverages pre-trained language models or trains lightweight LSTM models from scratch to generate synthetic text data.
 
-**Note**: The default model is `MOSTLY_AI/LSTMFromScratch-3m`, a lightweight LSTM model trained from scratch (GPU strongly recommended). You can also use pre-trained HuggingFace models by setting model to e.g. `microsoft/phi-1.5` (GPU required).
+**Note**: The default model is `MOSTLY_AI/LSTMFromScratch-3m`, a lightweight LSTM model trained from scratch (GPU strongly recommended). You can also use pretrained Hugging Face models (`model="<hub/repo>"`; GPU required). Verified checkpoints include `HuggingFaceTB/SmolLM2-135M`, `HuggingFaceTB/SmolLM3-3B`, `Qwen/Qwen3-0.6B`, and `microsoft/phi-4`.
 
 ### Model Training
 
@@ -309,5 +317,5 @@ lm.sample(
 
 Example notebooks demonstrating various use cases are available in the `examples` directory:
 - TabularARGN for flat tabular data [![Run on Colab](https://img.shields.io/badge/Open%20in-Colab-blue?logo=google-colab)](https://colab.research.google.com/github/mostly-ai/mostlyai-engine/blob/main/examples/flat.ipynb)
-- TabularARGN for sequential data [![Run on Colab](https://img.shields.io/badge/Open%20in-Colab-blue?logo=google-colab)](https://colab.research.google.com/github/mostly-ai/mostlyai-engine/blob/main/examples/language.ipynb)
-- LanguageModel for textual data [![Run on Colab](https://img.shields.io/badge/Open%20in-Colab-blue?logo=google-colab)](https://colab.research.google.com/github/mostly-ai/mostlyai-engine/blob/main/examples/sequential.ipynb)
+- TabularARGN for sequential data [![Run on Colab](https://img.shields.io/badge/Open%20in-Colab-blue?logo=google-colab)](https://colab.research.google.com/github/mostly-ai/mostlyai-engine/blob/main/examples/sequential.ipynb)
+- LanguageModel for textual data [![Run on Colab](https://img.shields.io/badge/Open%20in-Colab-blue?logo=google-colab)](https://colab.research.google.com/github/mostly-ai/mostlyai-engine/blob/main/examples/language.ipynb)

{mostlyai_engine-2.4.0 → mostlyai_engine-2.6.0}/mostlyai/engine/__init__.py

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

{mostlyai_engine-2.4.0 → mostlyai_engine-2.6.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.4.0 → mostlyai_engine-2.6.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.4.0 → mostlyai_engine-2.6.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.4.0 → mostlyai_engine-2.6.0}/mostlyai/engine/_language/common.py

@@ -52,8 +52,19 @@ def get_attention_implementation(config: PretrainedConfig) -> str | None:
 
 
 def load_base_model_and_config(
-    model_id_or_path: str | Path, device: torch.device, is_peft_adapter: bool, is_training: bool
+    model_id_or_path: str | Path,
+    device: torch.device,
+    is_peft_adapter: bool,
+    is_training: bool,
+    *,
+    differential_privacy: bool = False,
 ) -> tuple[PreTrainedModel, PretrainedConfig]:
+    """Load a HF base model (and config) for language training or inference.
+
+    When ``differential_privacy`` is True (Opacus DP training), the loader prefers
+    settings that keep per-sample gradients well-defined: float32 weights, no int4
+    training path, eager attention (not fused SDPA), and no gradient checkpointing.
+    """
     # opacus DP does not support parallel/sharded training
     model_id_or_path = str(model_id_or_path)
     if is_peft_adapter:
@@ -79,7 +90,13 @@ def load_base_model_and_config(
             "CUDA device was found but bitsandbytes is not available. Please use extra [gpu] to install bitsandbytes for quantization."
         )
     bf16_supported = is_bf16_supported(device)
-    if bf16_supported:
+    # Opacus needs reliable per-sample grads; bfloat16 params + grad_sample hooks are a poor match on many setups.
+    use_int4_training = is_gpu_training and is_bitsandbytes_available and not differential_privacy
+    if differential_privacy:
+        torch_dtype = torch.float32
+        # Eager attention keeps standard backward paths; fused SDPA can break Opacus grad_sample hooks.
+        attn_implementation = "eager"
+    elif bf16_supported:
         attn_implementation = get_attention_implementation(config)
         torch_dtype = torch.bfloat16
     else:
@@ -87,7 +104,7 @@ def load_base_model_and_config(
         torch_dtype = torch.float32
     if hasattr(config, "quantization_config"):
         quantization_config = AutoQuantizationConfig.from_dict(config.quantization_config)
-    elif is_gpu_training and is_bitsandbytes_available:
+    elif use_int4_training:
         quantization_config = BitsAndBytesConfig(
             load_in_4bit=True,
             bnb_4bit_quant_type="nf4",
@@ -123,8 +140,9 @@ def load_base_model_and_config(
     if isinstance(quantization_config, BitsAndBytesConfig):
         # convert all non-kbit layers to float32
         model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=False)
-    if is_gpu_training and model.supports_gradient_checkpointing:
+    if is_gpu_training and model.supports_gradient_checkpointing and not differential_privacy:
         # pay 50% time penalty for _large_ memory savings
+        # gradient checkpointing breaks Opacus per-sample gradient hooks
         _LOG.info("enable gradient checkpointing")
         model.gradient_checkpointing_enable()
         model.enable_input_require_grads()

{mostlyai_engine-2.4.0 → mostlyai_engine-2.6.0}/mostlyai/engine/_language/engine/hf_engine.py

@@ -22,11 +22,11 @@ import torch
 from peft import PeftModel
 from pydantic import BaseModel
 from transformers import AutoTokenizer
-from xgrammar.contrib.hf import LogitsProcessor
 
 from mostlyai.engine._language.common import load_base_model_and_config
 from mostlyai.engine._language.engine.base import EngineMetrics, LanguageEngine
 from mostlyai.engine._language.tokenizer_utils import tokenize_fn
+from mostlyai.engine._language.xgrammar_hf_logits import LogitsProcessor
 from mostlyai.engine._language.xgrammar_utils import create_compiled_grammars
 
 

{mostlyai_engine-2.4.0 → mostlyai_engine-2.6.0}/mostlyai/engine/_language/engine/vllm_engine.py

@@ -14,11 +14,6 @@
 
 from __future__ import annotations
 
-import os
-
-os.environ["VLLM_USE_V1"] = "1"
-
-
 import time
 from os import PathLike
 
@@ -28,7 +23,7 @@ from pydantic import BaseModel
 from transformers import AutoConfig, AutoTokenizer
 from vllm import LLM, SamplingParams
 from vllm.distributed import cleanup_dist_env_and_memory
-from vllm.inputs.data import TokensPrompt
+from vllm.inputs.llm import TokensPrompt
 from vllm.lora.request import LoRARequest
 from vllm.sampling_params import StructuredOutputsParams
 

{mostlyai_engine-2.4.0 → mostlyai_engine-2.6.0}/mostlyai/engine/_language/generation.py

@@ -150,7 +150,6 @@ def generate(
     _LOG.info("GENERATE_LANGUAGE started")
     t0_ = time.time()
     os.environ["VLLM_LOGGING_LEVEL"] = "WARNING"
-    os.environ["VLLM_NO_DEPRECATION_WARNING"] = "1"
 
     @contextlib.contextmanager
     def tqdm_disabled():

{mostlyai_engine-2.4.0 → mostlyai_engine-2.6.0}/mostlyai/engine/_language/interface.py

@@ -55,6 +55,9 @@ class LanguageModel(BaseEstimator):
         tgt_encoding_types: Dictionary mapping column names to encoding types.
             Example: {'category': 'LANGUAGE_CATEGORICAL', 'headline': 'LANGUAGE_TEXT'}
         model: The identifier of the language model to train. Defaults to MOSTLY_AI/LSTMFromScratch-3m.
+            Pretrained Hugging Face checkpoints are supported; verified examples include
+            HuggingFaceTB/SmolLM2-135M, HuggingFaceTB/SmolLM3-3B, Qwen/Qwen3-0.6B, and microsoft/phi-4
+            (GPU strongly recommended).
         max_training_time: Maximum training time in minutes. Defaults to 14400 (10 days).
         max_epochs: Maximum number of training epochs. Defaults to 100.
         batch_size: Per-device batch size for training and validation. If None, determined automatically.

{mostlyai_engine-2.4.0 → mostlyai_engine-2.6.0}/mostlyai/engine/_language/lstm.py

@@ -22,6 +22,17 @@ from transformers.modeling_outputs import CausalLMOutput
 _LOG = logging.getLogger(__name__)
 
 
+def _dplstm_state_dict_aliases(num_layers: int) -> dict[str, str]:
+    """Nested DPLSTM names -> flat ``nn.LSTM`` names (same storage; needed for checkpoint save)."""
+    aliases: dict[str, str] = {}
+    for i in range(num_layers):
+        aliases[f"lstm.l{i}.ih.weight"] = f"lstm.weight_ih_l{i}"
+        aliases[f"lstm.l{i}.ih.bias"] = f"lstm.bias_ih_l{i}"
+        aliases[f"lstm.l{i}.hh.weight"] = f"lstm.weight_hh_l{i}"
+        aliases[f"lstm.l{i}.hh.bias"] = f"lstm.bias_hh_l{i}"
+    return aliases
+
+
 class LSTMFromScratchConfig(PretrainedConfig):
     model_type = model_id = "MOSTLY_AI/LSTMFromScratch-3m"
 
@@ -54,7 +65,6 @@ class LSTMFromScratchLMHeadModel(PreTrainedModel, GenerationMixin):
 
     def __init__(self, config: LSTMFromScratchConfig):
         super().__init__(config)
-        self.config = config
 
         self.embedding = nn.Embedding(self.config.vocab_size, self.config.embedding_size)
         self.dropout = nn.Dropout(self.config.dropout)
@@ -77,6 +87,28 @@ class LSTMFromScratchLMHeadModel(PreTrainedModel, GenerationMixin):
         # this will be filled by left_to_right_padding() during the generation
         self.pad_token_id = None
 
+        # `_tied_weights_keys` is always a dict: empty unless DP (see `remove_tied_weights_from_state_dict`).
+        self._tied_weights_keys = _dplstm_state_dict_aliases(self.config.num_layers) if self.config.with_dp else {}
+
+        self.post_init()
+
+    def _init_weights(self, module: nn.Module) -> None:
+        # Keep PyTorch defaults for our main modules (historical behavior). HF post_init()
+        # still runs init_weights on the rest (e.g. any submodules inside DPLSTM).
+        if module in (self.embedding, self.lm_head, self.lstm):
+            return
+        super()._init_weights(module)
+
+    def get_expanded_tied_weights_keys(self, all_submodels: bool = False) -> dict[str, str]:
+        """
+        Transformers >= 5 sets `all_tied_weights_keys` from this. `self._tied_weights_keys` is also read when
+        saving (see `remove_tied_weights_from_state_dict`). Keep both in sync for DPLSTM aliases.
+        """
+        expanded = getattr(super(), "get_expanded_tied_weights_keys", None)
+        out: dict[str, str] = dict(expanded(all_submodels=all_submodels)) if expanded is not None else {}
+        out.update(self._tied_weights_keys)
+        return out
+
     def forward(
         self,
         input_ids: torch.Tensor,

{mostlyai_engine-2.4.0 → mostlyai_engine-2.6.0}/mostlyai/engine/_language/training.py

@@ -233,7 +233,8 @@ def _gpu_estimate_max_batch_size(
     model: PreTrainedModel | GradSampleModule, device: torch.device, max_tokens: int, initial_batch_size: int
 ) -> int:
     batch_size = 2 ** int(np.log2(initial_batch_size))
-    optimizer = torch.optim.AdamW(params=model.parameters())
+    # Match training optimizer: only trainable params (e.g. LoRA), for consistent memory probe.
+    optimizer = torch.optim.AdamW(params=[p for p in model.parameters() if p.requires_grad])
 
     # create test batch of zeros with estimated max sequence length
     def create_test_batch(batch_size: int):
@@ -339,7 +340,8 @@ def train(
         _LOG.info(f"{torch.cuda.device_count()=}")
         bf16_supported = is_bf16_supported(device)
         _LOG.info(f"{bf16_supported=}")
-        use_mixed_precision = bf16_supported and model != LSTMFromScratchConfig.model_id
+        use_mixed_precision = bf16_supported and model != LSTMFromScratchConfig.model_id and not with_dp
+        # DP uses float32 + no autocast (see load_base_model_and_config); bf16 autocast breaks Opacus grad_sample.
         _LOG.info(f"{use_mixed_precision=}")
 
         ctx_stats = workspace.ctx_stats.read()
@@ -450,7 +452,11 @@ def train(
             if resume_from_last_checkpoint:
                 tokenizer = AutoTokenizer.from_pretrained(model_id_or_path, **tokenizer_args)
                 model, _ = load_base_model_and_config(
-                    model_id_or_path, device=device, is_peft_adapter=False, is_training=True
+                    model_id_or_path,
+                    device=device,
+                    is_peft_adapter=False,
+                    is_training=True,
+                    differential_privacy=with_dp,
                 )
             else:
                 # fresh initialization of the custom tokenizer and LSTM model
@@ -468,6 +474,7 @@ def train(
                 device=device,
                 is_peft_adapter=resume_from_last_checkpoint,
                 is_training=True,
+                differential_privacy=with_dp,
             )
             tokenizer = AutoTokenizer.from_pretrained(model_id_or_path, **tokenizer_args)
             if tokenizer.eos_token is None:
@@ -584,7 +591,10 @@ def train(
             batch_size=val_batch_size,
             collate_fn=data_collator,
         )
-        optimizer = torch.optim.AdamW(params=model.parameters(), lr=initial_lr)
+        optimizer = torch.optim.AdamW(
+            params=[p for p in model.parameters() if p.requires_grad],
+            lr=initial_lr,
+        )  # frozen PEFT base weights must not be in Opacus optimizer (no grad_sample on unused params)
         early_stopper = EarlyStopper(val_loss_patience=4)
         lr_scheduler: LRScheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
             optimizer=optimizer,

mostlyai_engine-2.6.0/mostlyai/engine/_language/xgrammar_hf_logits.py

@@ -0,0 +1,69 @@
+# 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.
+
+"""xgrammar Hugging Face LogitsProcessor with fixes for our dependency stack."""
+
+from __future__ import annotations
+
+import torch
+import xgrammar as xgr
+from xgrammar.contrib.hf import LogitsProcessor as _UpstreamLogitsProcessor
+
+
+class LogitsProcessor(_UpstreamLogitsProcessor):
+    """
+    Same as ``xgrammar.contrib.hf.LogitsProcessor``, but passes a Python ``int`` to
+    ``GrammarMatcher.accept_token``. Upstream uses ``input_ids[i][-1]`` (a 0-dim tensor);
+    newer PyTorch / Transformers stacks leave that as ``torch.Tensor``, while xgrammar's
+    TVM FFI expects ``int``.
+    """
+
+    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
+        if len(self.matchers) == 0:
+            self.batch_size = input_ids.shape[0]
+            self.compiled_grammars = (
+                self.compiled_grammars if len(self.compiled_grammars) > 1 else self.compiled_grammars * self.batch_size
+            )
+            assert len(self.compiled_grammars) == self.batch_size, (
+                "The number of compiled grammars must be equal to the batch size."
+            )
+            self.matchers = [xgr.GrammarMatcher(self.compiled_grammars[i]) for i in range(self.batch_size)]
+            self.token_bitmask = xgr.allocate_token_bitmask(self.batch_size, self.full_vocab_size)
+
+        if input_ids.shape[0] != self.batch_size:
+            raise RuntimeError(
+                "Expect input_ids.shape[0] to be LogitsProcessor.batch_size."
+                + f"Got {input_ids.shape[0]} for the former, and {self.batch_size} for the latter."
+            )
+
+        if not self.prefilled:
+            self.prefilled = True
+        else:
+            for i in range(self.batch_size):
+                if not self.matchers[i].is_terminated():
+                    sampled_token = int(input_ids[i, -1].item())
+                    assert self.matchers[i].accept_token(sampled_token)
+
+        for i in range(self.batch_size):
+            if not self.matchers[i].is_terminated():
+                self.matchers[i].fill_next_token_bitmask(self.token_bitmask, i)
+
+        device_type = scores.device.type
+        if device_type != "cuda":
+            scores = scores.to("cpu")
+        xgr.apply_token_bitmask_inplace(scores, self.token_bitmask.to(scores.device))
+        if device_type != "cuda":
+            scores = scores.to(device_type)
+
+        return scores

{mostlyai_engine-2.4.0 → mostlyai_engine-2.6.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.4.0 → mostlyai_engine-2.6.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.4.0 → mostlyai_engine-2.6.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)
@@ -1067,7 +1066,7 @@ def generate(
                     sidx_df = encode_positional_column(sidx, max_seq_len=seq_steps, prefix=SIDX_SUB_COLUMN_PREFIX)
                     sidx_vals = {
                         c: torch.unsqueeze(
-                            torch.as_tensor(sidx_df[c].to_numpy(), device=model.device).type(torch.int),
+                            torch.as_tensor(sidx_df[c].to_numpy(copy=True), device=model.device).type(torch.int),
                             dim=-1,
                         )
                         for c in sidx_df
@@ -1111,7 +1110,7 @@ def generate(
                             sdec = pd.Series([0] * step_size)  # initial sequence index decile
                         sdec_vals = {
                             f"{SDEC_SUB_COLUMN_PREFIX}cat": torch.unsqueeze(
-                                torch.as_tensor(sdec.to_numpy(), device=model.device).type(torch.int), dim=-1
+                                torch.as_tensor(sdec.to_numpy(copy=True), device=model.device).type(torch.int), dim=-1
                             )
                         }
 
@@ -1120,7 +1119,9 @@ def generate(
                     if len(seed_step_encoded) > 0:
                         seed_vals = {
                             col: torch.unsqueeze(
-                                torch.as_tensor(seed_step_encoded[col].to_numpy(), device=model.device).type(torch.int),
+                                torch.as_tensor(seed_step_encoded[col].to_numpy(copy=True), device=model.device).type(
+                                    torch.int
+                                ),
                                 dim=-1,
                             )
                             for col in seed_step_encoded.columns
@@ -1224,7 +1225,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.4.0 → mostlyai_engine-2.6.0}/mostlyai/engine/training.py

@@ -47,7 +47,10 @@ def train(
     - `ModelStore`: Trained model checkpoints and logs.
 
     Args:
-        model: The identifier of the model to train. If tabular, defaults to MOSTLY_AI/Medium. If language, defaults to MOSTLY_AI/LSTMFromScratch-3m.
+        model: The identifier of the model to train. If tabular, defaults to MOSTLY_AI/Medium. If language,
+            defaults to MOSTLY_AI/LSTMFromScratch-3m. For language models, Hugging Face hub ids are supported;
+            verified pretrained checkpoints include HuggingFaceTB/SmolLM2-135M, HuggingFaceTB/SmolLM3-3B,
+            Qwen/Qwen3-0.6B, and microsoft/phi-4.
         max_training_time: Maximum training time in minutes. If None, defaults to 10 days.
         max_epochs: Maximum number of training epochs. If None, defaults to 100 epochs.
         batch_size: Per-device batch size for training and validation. If None, determined automatically.

{mostlyai_engine-2.4.0 → mostlyai_engine-2.6.0}/pyproject.toml

@@ -1,9 +1,9 @@
 [project]
 name = "mostlyai-engine"
-version = "2.4.0"
+version = "2.6.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,32 +24,32 @@ classifiers = [
 ]
 
 dependencies = [
-    "setuptools>=77.0.3",
+    "setuptools>=77.0.3,<81.0.0",  # vllm 0.20 caps setuptools for Python > 3.11
     "numpy>=2.0.0",
     "pandas>=2.2.0",
     "pyarrow>=16.0.0",
     "joblib>=1.4.2",
     "scikit-learn>=1.4.0",
     "psutil>=5.9.5,<6",  # upgrade when colab psutil is updated
-    "tokenizers>=0.21.0",
-    "transformers>=4.55.0",
+    "tokenizers>=0.21.1",  # vllm 0.20 lower bound
+    "transformers>=5.5.1",  # vllm 0.20 forbids 5.0–5.4 and 5.5.0; >=5.5.1 needed for newer HF model types (e.g. gemma4)
     "datasets>=3.0.0",
     "accelerate>=1.5.0",
-    "peft>=0.12.0",
+    "peft>=0.18.2",  # transformers 5.7+ checks min PEFT in model.add_adapter (integrations/peft.py)
     "huggingface-hub[hf-xet]>=0.30.2",
     "opacus>=1.5.4",
-    "xgrammar>=0.1.21",
+    "xgrammar>=0.1.32,<1.0.0",  # aligned with vllm 0.20
     "json-repair>=0.47.0",
-    "torch>=2.9.0,<2.10.0",
-    "torchaudio>=2.9.0,<2.10.0",
-    "torchvision>=0.24.0,<0.25.0"
+    "torch>=2.11.0,<2.12.0",
+    "torchaudio>=2.11.0,<2.12.0",
+    "torchvision>=0.26.0,<0.27.0"
 ]
 
 [project.optional-dependencies]
 gpu = [
     "bitsandbytes==0.42.0; sys_platform == 'darwin'",
     "bitsandbytes>=0.45.5; sys_platform == 'linux'",
-    "vllm==0.12.0; sys_platform == 'linux' or sys_platform == 'darwin'",
+    "vllm==0.20.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]