mostlyai-engine 1.1.12__tar.gz → 1.2.0__tar.gz

{mostlyai_engine-1.1.12 → mostlyai_engine-1.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mostlyai-engine
-Version: 1.1.12
+Version: 1.2.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
@@ -24,26 +24,25 @@ Classifier: Programming Language :: Python :: 3.13
 Classifier: Topic :: Software Development :: Libraries
 Classifier: Typing :: Typed
 Requires-Python: >=3.10
-Requires-Dist: accelerate>=0.31.0
-Requires-Dist: datasets<4,>=3.0.0
-Requires-Dist: formatron<0.5,>=0.4.11
-Requires-Dist: huggingface-hub<0.29,>=0.28.0
-Requires-Dist: joblib<2,>=1.4.2
+Requires-Dist: accelerate>=1.5.0
+Requires-Dist: datasets>=3.0.0
+Requires-Dist: huggingface-hub>=0.30.0
+Requires-Dist: joblib>=1.4.2
 Requires-Dist: json-repair>=0.30.0
-Requires-Dist: numpy<2,>=1.26.3
-Requires-Dist: opacus<2,>=1.5.2
+Requires-Dist: numpy>=1.26.3
+Requires-Dist: opacus>=1.5.2
 Requires-Dist: pandas~=2.2.0
 Requires-Dist: peft>=0.12.0
 Requires-Dist: psutil<6,>=5.9.5
 Requires-Dist: pyarrow>=16.0.0
-Requires-Dist: rich<14,>=13.9.4
 Requires-Dist: setuptools<76,>=75.0.0
-Requires-Dist: tokenizers<0.22,>=0.21.0
+Requires-Dist: tokenizers>=0.21.0
 Requires-Dist: torch==2.6.0; sys_platform == 'linux'
 Requires-Dist: torch>=2.6.0; sys_platform == 'darwin'
 Requires-Dist: torchaudio==2.6.0; sys_platform == 'linux'
 Requires-Dist: torchvision==0.21.0; sys_platform == 'linux'
-Requires-Dist: transformers<5,>=4.45.2
+Requires-Dist: transformers>=4.51.0
+Requires-Dist: xgrammar==0.1.17
 Provides-Extra: cpu
 Requires-Dist: torch==2.6.0+cpu; (sys_platform == 'linux') and extra == 'cpu'
 Requires-Dist: torch==2.6.0; (sys_platform != 'linux') and extra == 'cpu'

{mostlyai_engine-1.1.12 → mostlyai_engine-1.2.0}/mostlyai/engine/__init__.py

@@ -22,7 +22,7 @@ from mostlyai.engine.generation import generate
 
 
 __all__ = ["split", "analyze", "encode", "train", "generate", "init_logging"]
-__version__ = "1.1.12"
+__version__ = "1.2.0"
 
 # suppress specific warning related to os.fork() in multi-threaded processes
 warnings.filterwarnings("ignore", category=DeprecationWarning, message=".*multi-threaded.*fork.*")

{mostlyai_engine-1.1.12 → mostlyai_engine-1.2.0}/mostlyai/engine/_common.py

@@ -511,7 +511,7 @@ def skip_if_error(func: Callable) -> Callable:
 
 
 def encode_slen_sidx_sdec(vals: pd.Series, max_seq_len: int, prefix: str = "") -> pd.DataFrame:
-    assert vals.dtype == int or vals.dtype == "int64[pyarrow]"
+    assert is_integer_dtype(vals)
     if max_seq_len < SLEN_SIDX_DIGIT_ENCODING_THRESHOLD or prefix == SDEC_SUB_COLUMN_PREFIX:
         # encode slen and sidx as numeric_discrete
         df = pd.DataFrame({f"{prefix}cat": vals})

{mostlyai_engine-1.1.12 → mostlyai_engine-1.2.0}/mostlyai/engine/_language/engine/base.py

@@ -12,11 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from abc import ABC, abstractmethod
+from collections.abc import Generator
 from dataclasses import dataclass
-from abc import abstractmethod, ABC
-from collections.abc import Callable
 
-from formatron.formatter import FormatterBuilder
+from pydantic import BaseModel
 
 
 @dataclass
@@ -27,9 +27,7 @@ class EngineMetrics:
 
 class LanguageEngine(ABC):
     @abstractmethod
-    def initialize_logits_processors(
-        self, formatter_builders: list[FormatterBuilder], vocab_processors: list[Callable] | None = None
-    ):
+    def initialize_logits_processors(self, schemas: Generator[BaseModel]):
         pass
 
     @abstractmethod

{mostlyai_engine-1.1.12 → mostlyai_engine-1.2.0}/mostlyai/engine/_language/engine/hf_engine.py

@@ -12,22 +12,73 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from os import PathLike
+from __future__ import annotations
+
 import time
-from collections.abc import Callable
+from collections.abc import Generator
+from os import PathLike
 from pathlib import Path
 
 import torch
-from formatron.integrations.transformers import create_formatter_logits_processor_list
+import transformers
+import xgrammar as xgr
 from peft import PeftModel
-
+from pydantic import BaseModel
 from transformers import AutoTokenizer
+
 from mostlyai.engine._language.common import load_base_model_and_config
-from mostlyai.engine._language.formatron_utils import monkey_patch_formatron
+from mostlyai.engine._language.engine.base import EngineMetrics, LanguageEngine
 from mostlyai.engine._language.tokenizer_utils import tokenize_fn
+from mostlyai.engine._language.xgrammar_utils import create_compiled_grammars
+
+
+class XGrammarLogitsProcessor(transformers.LogitsProcessor):
+    """
+    Inspired by [LogitsProcessor](https://github.com/mlc-ai/xgrammar/blob/414473e7c029d0d9e2dfbeacb48afa946d0e3419/python/xgrammar/contrib/hf.py#L14).
+    HuggingFace's XGrammarLogitsProcessor cannot be reused. Logits processors must be initialized for each call to generate().
+    """
+
+    def __init__(self, compiled_grammars: list[xgr.CompiledGrammar]):
+        self.compiled_grammars = compiled_grammars
+        self.vocab_size = self.compiled_grammars[0].tokenizer_info.vocab_size
+        self.batch_size = len(compiled_grammars)
+
+        self.matchers: list[xgr.GrammarMatcher] | None = None
+        self.token_bitmask: torch.Tensor | None = None
+        self.prefilled = False
+
+    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
+        # lazily initialize GrammarMatchers and bitmask
+        if self.matchers is None:
+            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.vocab_size)
+
+        if input_ids.shape[0] != self.batch_size:
+            raise RuntimeError(
+                "Expect input_ids.shape[0] to be XGrammarLogitsProcessor.batch_size. "
+                + f"Got {input_ids.shape[0]} for the former, and {self.batch_size} for the latter."
+            )
+
+        if not self.prefilled:
+            # have not sampled a token yet
+            self.prefilled = True
+        else:
+            for i in range(self.batch_size):
+                if not self.matchers[i].is_terminated():
+                    sampled_token = input_ids[i][-1]
+                    assert self.matchers[i].accept_token(sampled_token)
 
-from mostlyai.engine._language.engine.base import EngineMetrics, LanguageEngine
-from formatron.formatter import FormatterBuilder
+        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
 
 
 class HuggingFaceEngine(LanguageEngine):
@@ -37,13 +88,13 @@ class HuggingFaceEngine(LanguageEngine):
         self.device = device
         self.max_new_tokens = max_new_tokens
         self.tokenizer_max_length = tokenizer_max_length
+        self.is_peft_adapter = (Path(model_path) / "adapter_config.json").exists()
 
-        is_peft_adapter = (Path(model_path) / "adapter_config.json").exists()
         model_path = str(model_path)
-        self._model, _ = load_base_model_and_config(
-            model_path, device=device, is_peft_adapter=is_peft_adapter, is_training=False
+        self._model, self._model_config = load_base_model_and_config(
+            model_path, device=device, is_peft_adapter=self.is_peft_adapter, is_training=False
         )
-        if is_peft_adapter:
+        if self.is_peft_adapter:
             self._model = PeftModel.from_pretrained(self._model, model_path, is_trainable=False)
             self._model = self._model.merge_and_unload()
             self._default_batch_size = 64
@@ -62,12 +113,10 @@ class HuggingFaceEngine(LanguageEngine):
         )
 
         # we can't enforce JSON output if LSTM tokenizer training was skipped
-        is_trained_lstm_tokenizer = not is_peft_adapter and self.tokenizer.vocab_size > len(
+        is_trained_lstm_tokenizer = not self.is_peft_adapter and self.tokenizer.vocab_size > len(
             self.tokenizer.special_tokens_map
         )
-        self._json_enforcing_possible = is_peft_adapter or is_trained_lstm_tokenizer
-        if self.supports_json_enforcing():
-            monkey_patch_formatron()
+        self._json_enforcing_possible = self.is_peft_adapter or is_trained_lstm_tokenizer
         self._logits_processors = None
 
     def get_default_batch_size(self) -> int:
@@ -76,12 +125,14 @@ class HuggingFaceEngine(LanguageEngine):
     def supports_json_enforcing(self) -> bool:
         return self._json_enforcing_possible
 
-    def initialize_logits_processors(
-        self, formatter_builders: list[FormatterBuilder], vocab_processors: list[Callable] | None = None
-    ):
-        self._logits_processors = create_formatter_logits_processor_list(
-            tokenizer=self.tokenizer, formatter_builders=formatter_builders, vocab_processors=vocab_processors
+    def initialize_logits_processors(self, schemas: Generator[BaseModel]):
+        compiled_grammars = create_compiled_grammars(
+            schemas=schemas,
+            tokenizer=self.tokenizer,
+            vocab_size=self._model_config.vocab_size,
+            is_peft_adapter=self.is_peft_adapter,
         )
+        self._logits_processors = [XGrammarLogitsProcessor(list(compiled_grammars))]
 
     def generate(
         self, text: list[str], sampling_temperature: float, sampling_top_p: float
@@ -107,22 +158,14 @@ class HuggingFaceEngine(LanguageEngine):
             temperature=sampling_temperature if do_sample else None,
             top_p=sampling_top_p if do_sample else None,
             bos_token_id=self.tokenizer.bos_token_id,
-            pad_token_id=self.tokenizer.eos_token_id,  # must be eos or will get ValueError from formatron
+            pad_token_id=self.tokenizer.pad_token_id,
             eos_token_id=self.tokenizer.eos_token_id,
         )
 
-        if self._logits_processors is not None:
-            # number of formatters must match the batch size, batch size is always reduced so this is fine
-            actual_batch_size = len(inputs["input_ids"])
-            self._logits_processors[0]._formatters = self._logits_processors[0]._formatters[:actual_batch_size]
-
         t_generate = time.time()
         outputs = self._model.generate(**inputs, **generate_kwargs, logits_processor=self._logits_processors)
         generate_time = time.time() - t_generate
 
-        if self._logits_processors:
-            self._logits_processors[0].reset()
-
         _, input_length = inputs["input_ids"].shape
         # truncate the prompt from the outputs
         outputs = outputs[:, input_length:]

{mostlyai_engine-1.1.12 → mostlyai_engine-1.2.0}/mostlyai/engine/_language/engine/vllm_engine.py

@@ -12,31 +12,30 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import contextlib
-from os import PathLike
+import gc
 import time
-import typing
+from collections.abc import Generator
+from os import PathLike
 
-from mostlyai.engine._language.common import is_bf16_supported
-from mostlyai.engine._language.engine.base import EngineMetrics, LanguageEngine
 import torch
-from formatron.config import EngineGenerationConfig
-from formatron.formatter import FormatterBuilder
-
+import xgrammar as xgr
 from peft import PeftConfig
-
-from transformers import AutoTokenizer, AutoConfig, PreTrainedTokenizerBase
-from mostlyai.engine._language.formatron_utils import monkey_patch_formatron
+from pydantic import BaseModel
+from transformers import AutoConfig, AutoTokenizer
 from vllm import LLM, SamplingParams
-from vllm.lora.request import LoRARequest
 from vllm.config import _get_and_verify_max_len
-from mostlyai.engine._language.tokenizer_utils import tokenize_fn
-from formatron.integrations.vllm import create_engine_vocabulary, FormattersLogitsProcessor
-from vllm.distributed import destroy_model_parallel, destroy_distributed_environment
-
-import gc
+from vllm.distributed import destroy_distributed_environment, destroy_model_parallel
+from vllm.lora.request import LoRARequest
 from vllm.platforms import current_platform
 
+from mostlyai.engine._language.common import is_bf16_supported
+from mostlyai.engine._language.engine.base import EngineMetrics, LanguageEngine
+from mostlyai.engine._language.tokenizer_utils import tokenize_fn
+from mostlyai.engine._language.xgrammar_utils import create_compiled_grammars
+
 
 def cleanup_dist_env_and_memory():
     """Copy from current main of vllm replace by import when possible"""
@@ -49,42 +48,76 @@ def cleanup_dist_env_and_memory():
         torch.cuda.empty_cache()
 
 
-def create_vllm_logits_processors(
-    llm: LLM,
-    formatter_builders: typing.Sequence[FormatterBuilder | None],
-    configs: typing.Sequence[EngineGenerationConfig] = None,
-    vocab_processors: list[typing.Callable] | None = None,
-) -> list[FormattersLogitsProcessor]:
+class XGrammarLogitsProcessor:
     """
-    Create list of formatter logits processors without creating engine vocabulary several times.
+    Inspired by [XGrammarLogitsProcessor](https://github.com/vllm-project/vllm/blob/a43aa183dc0cb2639044c15d272e0ce1941392b0/vllm/model_executor/guided_decoding/xgrammar_decoding.py#L280).
+    VLLM's XGrammarLogitsProcessor can be reused.
     """
-    tokenizer = llm.get_tokenizer()
-    vocab = create_engine_vocabulary(tokenizer, vocab_processors)
-    formatters = [
-        i.build(vocab, lambda tokens: tokenizer.decode(tokens)) if i is not None else None for i in formatter_builders
-    ]
-    return [FormattersLogitsProcessor([formatter], tokenizer.eos_token_id, configs) for formatter in formatters]
 
+    def __init__(self, compiled_grammar: xgr.CompiledGrammar):
+        self.compiled_grammar = compiled_grammar
+        self.tokenizer_info = compiled_grammar.tokenizer_info
+        self.batch_size = 1
 
-class MaskInvalidIndicesLogitsProcessor:
-    """
-    Certain models have output size greater than their vocabulary size.
-    This logits processor masks the output indices that do not correspond
-    to a token id.
-    """
+        self.matchers: list[xgr.GrammarMatcher] | None = None
+        self.token_bitmask: torch.Tensor | None = None
+        self.prefilled = False
 
-    def __init__(self, tokenizer: PreTrainedTokenizerBase):  # : PretrainedTokenizer
-        self.mask: torch.Tensor | None = None
-        self.valid_token_ids = torch.tensor(list(tokenizer.vocab.values()))
+    def __call__(self, input_ids: tuple[int], scores: torch.Tensor) -> torch.Tensor:
+        # lazily initialize GrammarMatchers and bitmask
+        if self.matchers is None:
+            self.matchers = [xgr.GrammarMatcher(self.compiled_grammar) for _ in range(self.batch_size)]
+            self.token_bitmask = xgr.allocate_token_bitmask(self.batch_size, self.tokenizer_info.vocab_size)
 
-    def __call__(self, input_ids: list[int], scores: torch.Tensor) -> torch.Tensor:
-        if self.mask is None:
-            ninf = float("-inf")
-            self.mask = torch.full_like(scores, ninf)
-            self.mask[..., self.valid_token_ids] = 0
-        scores = scores + self.mask
+        if not self.prefilled:
+            # have not sampled a token yet
+            self.prefilled = True
+        else:
+            for i, matcher in enumerate(self.matchers):
+                if not matcher.is_terminated():
+                    sampled_token = input_ids[-1]
+                    assert self.matchers[i].accept_token(sampled_token)
+
+        for i, matcher in enumerate(self.matchers):
+            if not matcher.is_terminated():
+                matcher.fill_next_token_bitmask(self.token_bitmask, i)
+
+        # token_bitmask is a CPU tensor for use with accept_token and
+        # fill_next_token_bitmask so we move it to the device of scores
+        device_type = scores.device.type
+        dtype = scores.dtype
+        if device_type != "cuda":
+            # xgrammar on cpu only supports float32 scores
+            # see: https://github.com/mlc-ai/xgrammar/blob/c1b64920cad24f44f235778c1c00bb52d57da01a/python/xgrammar/kernels/apply_token_bitmask_inplace_cpu.py#L22
+            scores = scores.to("cpu").float().unsqueeze(0)
+
+        # Note: In this method, if the tensors have different dimensions
+        # on CPU device fails, but on GPU it runs without error. Hence the
+        # unsqueeze above for scores, to match the token bitmask shape
+        xgr.apply_token_bitmask_inplace(scores, self.token_bitmask.to(scores.device, non_blocking=True))
+        if device_type != "cuda":
+            scores = scores.to(dtype).to(device_type).squeeze()
         return scores
 
+    def clone(self) -> XGrammarLogitsProcessor:
+        """Create a new instance with shared compiled grammar but separate state"""
+        new_processor = XGrammarLogitsProcessor(self.compiled_grammar)
+
+        # create fresh matchers for the new sequence
+        if self.matchers is not None:
+            new_processor.matchers = [xgr.GrammarMatcher(self.compiled_grammar) for _ in range(self.batch_size)]
+
+        # create a new token bitmask with the same size
+        if self.token_bitmask is not None:
+            new_processor.token_bitmask = self.token_bitmask
+
+        new_processor.batch_size = self.batch_size
+
+        # reset prefilled state for new sequence
+        new_processor.prefilled = False
+
+        return new_processor
+
 
 class VLLMEngine(LanguageEngine):
     def __init__(
@@ -113,7 +146,6 @@ class VLLMEngine(LanguageEngine):
             swap_space=0,
             disable_log_stats=True,
         )
-        self._base_logits_processors = [MaskInvalidIndicesLogitsProcessor(self.llm.get_tokenizer())]
         self.tokenizer = AutoTokenizer.from_pretrained(
             model_path,
             padding_side="left",
@@ -124,7 +156,6 @@ class VLLMEngine(LanguageEngine):
             add_eos_token=False,
         )
         self._logits_processors = None
-        monkey_patch_formatron()
 
     def get_default_batch_size(self) -> int:
         return 192
@@ -132,12 +163,14 @@ class VLLMEngine(LanguageEngine):
     def supports_json_enforcing(self) -> bool:
         return True
 
-    def initialize_logits_processors(
-        self, formatter_builders: list[FormatterBuilder], vocab_processors: list[typing.Callable] | None = None
-    ):
-        self._logits_processors = create_vllm_logits_processors(
-            llm=self.llm, formatter_builders=formatter_builders, configs=None, vocab_processors=vocab_processors
+    def initialize_logits_processors(self, schemas: Generator[BaseModel]):
+        compiled_grammars = create_compiled_grammars(
+            schemas=schemas,
+            tokenizer=self.llm.get_tokenizer(),
+            vocab_size=self.llm.llm_engine.get_model_config().get_vocab_size(),
+            is_peft_adapter=True,
         )
+        self._logits_processors = [XGrammarLogitsProcessor(compiled_grammar) for compiled_grammar in compiled_grammars]
 
     def generate(
         self, text: list[str], sampling_temperature: float, sampling_top_p: float
@@ -156,23 +189,15 @@ class VLLMEngine(LanguageEngine):
         tokenize_time = time.time() - t_tokenize
 
         actual_batch_size = len(inputs["input_ids"])
-        if self._logits_processors is not None:
-            sampling_params = [
-                SamplingParams(
-                    max_tokens=self.max_new_tokens,
-                    temperature=sampling_temperature,
-                    top_p=sampling_top_p,
-                    logits_processors=[lp, *self._base_logits_processors],
-                )
-                for lp in self._logits_processors[:actual_batch_size]
-            ]
-        else:
-            sampling_params = SamplingParams(
+        sampling_params = [
+            SamplingParams(
                 max_tokens=self.max_new_tokens,
                 temperature=sampling_temperature,
                 top_p=sampling_top_p,
-                logits_processors=self._base_logits_processors,
+                logits_processors=[lp],
             )
+            for lp in self._logits_processors[:actual_batch_size]
+        ]
         t_generate = time.time()
         outputs = self.llm.generate(
             prompts=None,
@@ -182,9 +207,6 @@ class VLLMEngine(LanguageEngine):
             lora_request=self._lora_request,
         )
         generate_time = time.time() - t_generate
-        if self._logits_processors is not None:
-            for lp in self._logits_processors:
-                lp.reset()
         metrics = EngineMetrics(tokenize_time=tokenize_time, generate_time=generate_time)
         return [r.outputs[0].token_ids for r in outputs], metrics
 

{mostlyai_engine-1.1.12 → mostlyai_engine-1.2.0}/mostlyai/engine/_language/generation.py

@@ -40,12 +40,9 @@ from mostlyai.engine._encoding_types.language.numeric import decode_language_num
 from mostlyai.engine._encoding_types.language.text import decode_text
 from mostlyai.engine._language.common import MAX_LENGTH
 from mostlyai.engine._language.encoding import encode_df
+from mostlyai.engine._language.xgrammar_utils import create_schemas
 from mostlyai.engine._workspace import ensure_workspace_dir, Workspace, reset_dir
-from mostlyai.engine._language.formatron_utils import (
-    get_formatter_builders,
-    prepare_seed_for_formatron,
-    get_vocab_processors,
-)
+from mostlyai.engine._language.xgrammar_utils import ensure_seed_can_be_tokenized
 from mostlyai.engine.domain import ModelEncodingType, RareCategoryReplacementMethod
 
 INVALID_VALUE = "_INVALID_"  # when JSON parsing fails, the values of target columns will be set to this
@@ -274,8 +271,7 @@ def generate(
             batch_size = sample_size
         _LOG.info(f"{batch_size=}")
 
-        # prepare seed data for clean consumption by formatron
-        seed_data = prepare_seed_for_formatron(seed_data, engine.tokenizer)
+        seed_data = ensure_seed_can_be_tokenized(seed_data, engine.tokenizer)
         seeded_tgt_columns = seed_data.columns.to_list()
 
         total_tokenize_fn_time = 0
@@ -284,14 +280,16 @@ def generate(
 
         enforce_json_output = engine.supports_json_enforcing()
         _LOG.info(f"{enforce_json_output=}")
-        formatron_vocab_processors = get_vocab_processors(is_peft_adapter)
 
-        if enforce_json_output and len(seeded_tgt_columns) == 0:
+        initialize_logits_processors_once = len(seeded_tgt_columns) == 0 and engine.__class__.__name__ == "VLLMEngine"
+        if enforce_json_output and initialize_logits_processors_once:
             t0 = time.time()
-            formatter_builders = get_formatter_builders(
-                size=batch_size, stats=tgt_stats, rare_category_replacement_method=rare_category_replacement_method
+            schemas = create_schemas(
+                size=batch_size,
+                stats=tgt_stats,
+                rare_category_replacement_method=rare_category_replacement_method,
             )
-            engine.initialize_logits_processors(formatter_builders, formatron_vocab_processors)
+            engine.initialize_logits_processors(schemas=schemas)
             total_logits_processor_build_time += time.time() - t0
 
         # keep at most 500k samples in memory before decoding and writing to disk
@@ -305,15 +303,14 @@ def generate(
             ctx_batch = ctx_data.iloc[samples_processed : samples_processed + batch_size]
             ctx_keys = ctx_batch[ctx_primary_key]
 
-            if enforce_json_output and len(seeded_tgt_columns) > 0:
+            if enforce_json_output and not initialize_logits_processors_once:
                 t0 = time.time()
-                # some columns are seeded, so we need to create a new logits processor for each batch
-                formatter_builders = get_formatter_builders(
+                schemas = create_schemas(
                     seed_df=sample_seed_batch,
                     stats=tgt_stats,
                     rare_category_replacement_method=rare_category_replacement_method,
                 )
-                engine.initialize_logits_processors(formatter_builders, formatron_vocab_processors)
+                engine.initialize_logits_processors(schemas=schemas)
                 total_logits_processor_build_time += time.time() - t0
 
             outputs, metrics = engine.generate(

mostlyai_engine-1.2.0/mostlyai/engine/_language/xgrammar_utils.py

@@ -0,0 +1,147 @@
+# 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.
+
+from __future__ import annotations
+
+import json
+from collections.abc import Generator
+from typing import Literal
+
+import pandas as pd
+import xgrammar as xgr
+from pydantic import BaseModel, Field, SkipValidation, create_model
+from transformers import PreTrainedTokenizerBase
+from xgrammar.testing import _json_schema_to_ebnf
+
+from mostlyai.engine._common import STRING
+from mostlyai.engine._encoding_types.language.categorical import CATEGORICAL_UNKNOWN_TOKEN
+from mostlyai.engine.domain import ModelEncodingType, RareCategoryReplacementMethod
+
+JSON_NULL = "null"
+
+
+def prepend_grammar_root_with_space(grammar: str) -> str:
+    # XGrammar always starts with "{" when enforcing JSON Schema
+    # training is done on strings like ` {...} {...}`
+    # later, generation is done with prompt like ` {...}`
+    # so the first output token must be a space
+    start_of_grammar = 'root ::= "{"'
+    start_of_grammar_with_space = 'root ::= " {"'
+    assert start_of_grammar in grammar
+    return grammar.replace(start_of_grammar, start_of_grammar_with_space)
+
+
+def ensure_seed_can_be_tokenized(sample_seed: pd.DataFrame, tokenizer: PreTrainedTokenizerBase) -> pd.DataFrame:
+    def transform(x: str | pd._libs.missing.NAType) -> str:
+        if pd.isna(x):
+            null = tokenizer.decode(tokenizer.encode(JSON_NULL), skip_special_tokens=True)
+            # xgrammar needs to be able to express JSON_NULL with available vocabulary
+            # if that's the case, harmonize null-like values to None (e.g. pd.NA would cause xgrammar to fail)
+            # otherwise, fallback to empty string
+            return None if null == JSON_NULL else ""
+        # skip tokens unseen during training
+        return tokenizer.decode(tokenizer.encode(x), skip_special_tokens=True)
+
+    return sample_seed.astype(STRING).map(transform)
+
+
+def create_schemas(
+    *,
+    seed_df: pd.DataFrame | None = None,
+    size: int | None = None,
+    stats: dict,
+    rare_category_replacement_method: RareCategoryReplacementMethod,
+) -> Generator[BaseModel]:
+    assert (seed_df is not None) ^ (size is not None), "exactly one of seed_df or size must be provided"
+    if seed_df is None:
+        seed_df = pd.DataFrame(index=range(size))
+    unseeded_fields = [c for c in list(stats["columns"].keys()) if c not in seed_df.columns.to_list()]
+    field_types = {
+        t: [col for col, col_stats in stats["columns"].items() if col_stats["encoding_type"] == t]
+        for t in ModelEncodingType
+    }
+    categorical_fields = field_types.get(ModelEncodingType.language_categorical, [])
+    numeric_fields = field_types.get(ModelEncodingType.language_numeric, [])
+    datetime_fields = field_types.get(ModelEncodingType.language_datetime, [])
+    cache = {}
+    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()])))
+        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]
+        for field_name in unseeded_fields:
+            if field_name in categorical_fields:
+                categories = stats["columns"][field_name]["categories"]
+                if rare_category_replacement_method == RareCategoryReplacementMethod.sample and len(categories) > 1:
+                    categories = [c for c in categories if c != CATEGORICAL_UNKNOWN_TOKEN]
+                model_dict[field_name] = (Literal[tuple(categories)], ...)  # type: ignore[valid-type]
+            elif field_name in numeric_fields:
+                max_scale = stats["columns"][field_name]["max_scale"]
+                min_min5 = min(stats["columns"][field_name]["min5"])
+                max_max5 = max(stats["columns"][field_name]["max5"])
+                if max_scale == 0:
+                    model_dict[field_name] = (SkipValidation[int], Field(ge=min_min5, le=max_max5))
+                else:
+                    model_dict[field_name] = (
+                        SkipValidation[float],
+                        Field(ge=min_min5, le=max_max5, decimal_places=max_scale),
+                    )
+            elif field_name in datetime_fields:
+                model_dict[field_name] = (
+                    SkipValidation[str],
+                    Field(
+                        pattern=r"""(19\d{2}|20\d{2})-(0[1-9]|1[0-2])-(0[1-9]|1[0-9]|2[0-9]|3[0-1])T([0-1][0-9]|2[0-3]):([0-5][0-9]):([0-5][0-9])"""
+                    ),
+                )
+            else:
+                model_dict[field_name] = (str, ...)
+        schema = create_model("TargetModel", **model_dict)
+        cache[cache_key] = schema
+        yield schema
+
+
+def _get_tokenizer_info_for_lstm(tokenizer: PreTrainedTokenizerBase, vocab_size: int):
+    # trimmed down version of xgr.TokenizerInfo.from_huggingface
+    # the original function sets vocab_type to VocabType.RAW,
+    # but LSTM tokenizer needs VocabType.BYTE_FALLBACK, because of the usage of metaspace ("▁")
+    encoded_vocab = [""] * vocab_size
+    for token, idx in tokenizer.get_vocab().items():
+        if idx < vocab_size:
+            encoded_vocab[idx] = token
+    tokenizer_info = xgr.TokenizerInfo(
+        encoded_vocab,
+        vocab_type=xgr.VocabType.BYTE_FALLBACK,
+        vocab_size=vocab_size,
+        stop_token_ids=[tokenizer.eos_token_id],
+        add_prefix_space=True,
+    )
+    return tokenizer_info
+
+
+def create_compiled_grammars(
+    schemas: Generator[BaseModel], tokenizer: PreTrainedTokenizerBase, vocab_size: int, is_peft_adapter: bool
+) -> Generator[xgr.CompiledGrammar]:
+    # in general, there might be misalignment between the model's and tokenizer's vocab_size
+    # the former is expected by XGrammar
+    make_tokenizer_info = xgr.TokenizerInfo.from_huggingface if is_peft_adapter else _get_tokenizer_info_for_lstm
+    tokenizer_info = make_tokenizer_info(tokenizer, vocab_size=vocab_size)
+    grammar_compiler = xgr.GrammarCompiler(tokenizer_info)
+    schemas = (json.dumps(schema.model_json_schema()) for schema in schemas)
+    grammars = (_json_schema_to_ebnf(schema) for schema in schemas)
+    grammars = (prepend_grammar_root_with_space(grammar) for grammar in grammars)
+    compiled_grammars = (grammar_compiler.compile_grammar(grammar) for grammar in grammars)
+    return compiled_grammars

{mostlyai_engine-1.1.12 → mostlyai_engine-1.2.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "mostlyai-engine"
-version = "1.1.12"
+version = "1.2.0"
 description = "Synthetic Data Engine"
 authors = [{ name = "MOSTLY AI", email = "dev@mostly.ai" }]
 requires-python = ">=3.10"
@@ -26,21 +26,20 @@ classifiers = [
 
 dependencies = [
     "setuptools>=75.0.0,<76",
-    "numpy>=1.26.3,<2",
+    "numpy>=1.26.3",
     "pandas~=2.2.0",
     "pyarrow>=16.0.0",
-    "joblib>=1.4.2,<2",
+    "joblib>=1.4.2",
     "psutil>=5.9.5,<6",  # upgrade when colab psutil is updated
-    "tokenizers>=0.21.0,<0.22",
-    "transformers>=4.45.2,<5",
-    "datasets>=3.0.0,<4",
-    "accelerate>=0.31.0",
+    "tokenizers>=0.21.0",
+    "transformers>=4.51.0",
+    "datasets>=3.0.0",
+    "accelerate>=1.5.0",
     "peft>=0.12.0",
-    "huggingface-hub>=0.28.0,<0.29",
-    "opacus>=1.5.2,<2",
-    "formatron>=0.4.11,<0.5",
+    "huggingface-hub>=0.30.0",
+    "opacus>=1.5.2",  # switch to 1.5.4 (once released) to allow numpy 2
+    "xgrammar==0.1.17",  # for vllm 0.8.2 compatibility
     "json-repair>=0.30.0",
-    "rich>=13.9.4,<14",
     "torch>=2.6.0; sys_platform == 'darwin'",
     "torch==2.6.0; sys_platform == 'linux'",  # for vllm 0.8.2 compatibility
     "torchaudio==2.6.0; sys_platform == 'linux'",  # for vllm 0.8.2 compatibility
@@ -118,8 +117,11 @@ include = ["mostlyai/engine"]
 [tool.hatch.build.targets.wheel]
 include = ["mostlyai/engine"]
 
+[tool.hatch.metadata]
+allow-direct-references = true
+
 [build-system]
-requires = ["hatchling"]
+requires = ["hatchling", "hatch-vcs"]
 build-backend = "hatchling.build"
 
 [tool.ruff]

mostlyai_engine-1.1.12/mostlyai/engine/_language/formatron_utils.py

@@ -1,264 +0,0 @@
-# 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.
-
-
-import typing
-
-
-import pandas as pd
-from formatron.schemas.pydantic import ClassSchema
-from json import JSONDecodeError
-from pydantic import Field, SkipValidation, ValidationError
-from formatron.formatter import FormatterBuilder
-from formatron import schemas
-from formatron.formats import json
-from typing import Literal
-from pydantic import create_model
-from transformers import PreTrainedTokenizerBase
-from mostlyai.engine._encoding_types.language.categorical import CATEGORICAL_UNKNOWN_TOKEN
-from mostlyai.engine.domain import ModelEncodingType, RareCategoryReplacementMethod
-
-JSON_NULL = "null"
-
-
-def prepare_seed_for_formatron(sample_seed: pd.DataFrame, tokenizer: PreTrainedTokenizerBase) -> pd.DataFrame:
-    def transform(x: str | None) -> str:
-        if pd.isna(x):
-            null = tokenizer.decode(tokenizer.encode(JSON_NULL), skip_special_tokens=True)
-            # formatron needs to be able to express JSON_NULL with available vocabulary
-            # if that's the case, harmonize null-like values to None (e.g. pd.NA would cause formatron to fail)
-            # otherwise, fallback to empty string
-            return None if null == JSON_NULL else ""
-        # skip tokens unseen during training
-        return tokenizer.decode(tokenizer.encode(x), skip_special_tokens=True)
-
-    return sample_seed.astype("string[pyarrow]").map(transform)
-
-
-def get_formatter_builders(
-    *,
-    seed_df: pd.DataFrame | None = None,
-    size: int | None = None,
-    stats: dict,
-    rare_category_replacement_method: RareCategoryReplacementMethod,
-) -> list[FormatterBuilder]:
-    assert (seed_df is not None) ^ (size is not None), "exactly one of seed_df or size must be provided"
-    formatter_builders = []
-    if seed_df is None:
-        seed_df = pd.DataFrame(index=range(size))
-    unseeded_fields = [c for c in list(stats["columns"].keys()) if c not in seed_df.columns.to_list()]
-    field_types = {
-        t: [col for col, col_stats in stats["columns"].items() if col_stats["encoding_type"] == t]
-        for t in ModelEncodingType
-    }
-    categorical_fields = field_types.get(ModelEncodingType.language_categorical, [])
-    numeric_fields = field_types.get(ModelEncodingType.language_numeric, [])
-    datetime_fields = field_types.get(ModelEncodingType.language_datetime, [])
-    cache = {}
-    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()])))
-        if cache_key in cache:
-            formatter_builders.append(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]
-        for field_name in unseeded_fields:
-            if field_name in categorical_fields:
-                categories = stats["columns"][field_name]["categories"]
-                if rare_category_replacement_method == RareCategoryReplacementMethod.sample and len(categories) > 1:
-                    categories = [c for c in categories if c != CATEGORICAL_UNKNOWN_TOKEN]
-                model_dict[field_name] = (Literal[tuple(categories)], ...)  # type: ignore[valid-type]
-            elif field_name in numeric_fields:
-                max_scale = stats["columns"][field_name]["max_scale"]
-                min_min5 = min(stats["columns"][field_name]["min5"])
-                max_max5 = max(stats["columns"][field_name]["max5"])
-                if max_scale == 0:
-                    model_dict[field_name] = (SkipValidation[int], Field(ge=min_min5, le=max_max5))
-                else:
-                    model_dict[field_name] = (
-                        SkipValidation[float],
-                        Field(ge=min_min5, le=max_max5, decimal_places=max_scale),
-                    )
-            elif field_name in datetime_fields:
-                model_dict[field_name] = (
-                    SkipValidation[str],
-                    Field(
-                        pattern=r"""(19\\d{2}|20\\d{2})-(0[1-9]|1[0-2])-(0[1-9]|1[0-9]|2[0-9]|3[0-1])T([0-1][0-9]|2[0-3]):([0-5][0-9]):([0-5][0-9])"""
-                    ),
-                )
-            else:
-                model_dict[field_name] = (str, ...)
-        schema = create_model("TargetModel", **model_dict, __base__=MostlyClassSchema)
-        formatter_builder = FormatterBuilder()
-        formatter_builder.append_str(f"{formatter_builder.json(schema, capture_name=None)}")
-        cache[cache_key] = formatter_builder
-        formatter_builders.append(formatter_builder)
-    return formatter_builders
-
-
-def get_vocab_processors(is_peft_adapter: bool) -> list[typing.Callable] | None:
-    if not is_peft_adapter:
-
-        def update_vocab_lstm(token_to_char: dict[bytes, bytes]):
-            """
-            Maps special tokens ("▁", "␊") back to their original representation (" ", "\n")
-            (used in LSTM tokenizer)
-            """
-            token_to_char["\u2581".encode()] = b" "  # "▁" -> " "
-            token_to_char["\u240a".encode()] = b"\n"  # "␊" -> "\n"
-
-        return [update_vocab_lstm]
-    return None
-
-
-class MostlyClassSchema(ClassSchema):
-    @classmethod
-    def from_json(cls, _json: str) -> "MostlyClassSchema":
-        """
-        Create a MostlyClassSchema from a JSON string.
-        """
-        try:
-            return cls.model_validate_json(_json)
-        except ValidationError as e:
-            for error in e.errors():
-                if error["type"] == "json_invalid":
-                    raise JSONDecodeError(
-                        f"Caught pydantic ValidationError {e}, reraising as JSONDecodeError", _json, 0
-                    )
-            raise e
-
-
-# copy formatron: direct copy from formatron
-def _string_metadata(current: type, nonterminal: str):
-    min_length = current.metadata.get("min_length")
-    max_length = current.metadata.get("max_length")
-    pattern = current.metadata.get("pattern")
-    substring_of = current.metadata.get("substring_of")
-    if pattern:
-        assert not (min_length or max_length or substring_of), (
-            "pattern is mutually exclusive with min_length, max_length and substring_of"
-        )
-    if substring_of:
-        assert not (min_length or max_length or pattern), (
-            "substring_of is mutually exclusive with min_length, max_length and pattern"
-        )
-    repetition_map = {
-        (True, False): f"{{{min_length},}}",
-        (False, True): f"{{0,{max_length}}}",
-        (True, True): f"{{{min_length},{max_length}}}",
-    }
-    repetition = repetition_map.get((min_length is not None, max_length is not None))
-    if repetition is not None:
-        return (
-            rf"""{nonterminal} ::= #'"([^\\\\"\u0000-\u001f]|\\\\["\\\\bfnrt/]|\\\\u[0-9A-Fa-f]{{4}}){repetition}"';
-""",
-            [],
-        )
-    if pattern is not None:
-        pattern = pattern.replace("'", "\\'")
-        return f"""{nonterminal} ::= #'"{pattern}"';\n""", []
-    if substring_of is not None:
-        return f"""{nonterminal} ::= '"' #substrs{repr(substring_of)} '"';\n""", []
-
-
-# completely altered vs formatron
-def _number_metadata(current: type, nonterminal: str):
-    # For now only constrains number of digits and whether it is negative
-    gt = current.metadata.get("gt")
-    ge = current.metadata.get("ge")
-    lt = current.metadata.get("lt")
-    le = current.metadata.get("le")
-    if lt is not None or gt is not None:
-        raise NotImplementedError("gt and lt are not supported for number metadata")
-    if le < ge:
-        raise ValueError("le must be greater than or equal to ge")
-
-    pattern_parts = []
-    if issubclass(current.type, float):
-        le, le_frac = str(le).split(".")
-        ge, ge_frac = str(ge).split(".")
-        le, le_frac = int(le), int(le_frac)
-        ge, ge_frac = int(ge), int(ge_frac)
-        decimal_places = current.metadata.get("decimal_places")
-
-    if ge is not None and le is not None:
-        if ge < 0 and le < 0:
-            pattern_parts.append("-")
-            min_num = abs(le)
-            max_num = abs(ge)
-            max_digits = len(str(max_num))
-            min_digits = len(str(min_num))
-            pattern_parts.append(rf"([1-9][0-9]{{{min_digits - 1},{max_digits - 1}}})")
-        elif ge > 0:
-            min_num = ge
-            max_num = le
-            max_digits = len(str(max_num))
-            min_digits = len(str(min_num))
-            pattern_parts.append(rf"([1-9][0-9]{{{min_digits - 1},{max_digits - 1}}})")
-        else:
-            if ge < 0:
-                pattern_parts.append("-?")
-            max_digits = max(len(str(abs(ge))), len(str(abs(le))))
-            pattern_parts.append(rf"(0|[1-9][0-9]{{0,{max_digits - 1}}})")
-
-    if issubclass(current.type, float):
-        pattern_parts.append(rf"(\\.[0-9]{{0,{decimal_places}}})?")
-
-    pattern = "".join(pattern_parts)
-    return f"""{nonterminal} ::= #"{pattern}";\n""", []
-
-
-# copy formatron: removed sequence metadata since unnecessary and altered number_metadata to use ours
-def _metadata(current: type, nonterminal: str):
-    if isinstance(current, schemas.schema.TypeWithMetadata):
-        original = typing.get_origin(current.type)
-        if original is None:
-            original = current.type
-        if not current.metadata:
-            return "", [(current.type, nonterminal)]
-        if isinstance(current.type, type) and issubclass(current.type, str):
-            return _string_metadata(current, nonterminal)
-        elif isinstance(current.type, type) and issubclass(current.type, (int, float)):
-            return _number_metadata(current, nonterminal)
-    return None
-
-
-def monkey_patch_formatron():
-    FORMATRON_WHITESPACE_MAX_REPETITIONS = 10
-    SPACE_NONTERMINAL = f"[ \t\n\r]{{0,{FORMATRON_WHITESPACE_MAX_REPETITIONS}}}"
-
-    # Copy from formatron, altered to have limited whitespace repetitions and datetime format
-    json.GRAMMAR_HEADER = rf"""integer ::= #"-?(0|[1-9]\\d*)";
-    number ::= #"-?(0|[1-9]\\d*)(\\.\\d+)?([eE][+-]?\\d+)?";
-    string ::= #'"([^\\\\"\u0000-\u001f]|\\\\["\\\\bfnrt/]|\\\\u[0-9A-Fa-f]{{4}})*"';
-    boolean ::= "true"|"false";
-    null ::= "null";
-    array ::= array_begin (json_value (comma json_value)*)? array_end;
-    object ::= object_begin (string colon json_value (comma string colon json_value)*)? object_end;
-    json_value ::= number|string|boolean|null|array|object;
-    comma ::= #"{SPACE_NONTERMINAL},{SPACE_NONTERMINAL}";
-    colon ::= #"{SPACE_NONTERMINAL}:{SPACE_NONTERMINAL}";
-    object_begin ::= #" \\{{{SPACE_NONTERMINAL}";
-    object_end ::= #"{SPACE_NONTERMINAL}\\}}";
-    array_begin ::= #"\\[{SPACE_NONTERMINAL}";
-    array_end ::= #"{SPACE_NONTERMINAL}\\]";
-    """
-
-    def alter_type_to_nonterminals_metadata_inplace(type_to_nonterminals: list[typing.Callable]):
-        metadata_idx = [idx for idx, fn in enumerate(type_to_nonterminals) if fn.__name__ == "metadata"]
-        if len(metadata_idx) == 1:
-            type_to_nonterminals[metadata_idx[0]] = _metadata
-
-    alter_type_to_nonterminals_metadata_inplace(json._type_to_nonterminals)