mostlyai-engine 1.0.0__py3-none-any.whl

mostlyai/engine/_language/encoding.py

@@ -0,0 +1,162 @@
+# 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 json
+import logging
+import time
+import warnings
+from pathlib import Path
+
+import pandas as pd
+from pandas.core.dtypes.common import is_datetime64_any_dtype
+from tokenizers.pre_tokenizers import ByteLevel
+
+from mostlyai.engine._common import is_sequential, ProgressCallback, ProgressCallbackWrapper, TABLE_COLUMN_INFIX
+from mostlyai.engine._workspace import ensure_workspace_dir, Workspace, reset_dir
+
+_LOG = logging.getLogger(__name__)
+
+
+def format_df(df: pd.DataFrame, columns: list[str], is_target: bool = False) -> pd.DataFrame:
+    df = df[columns].copy()
+    # Some columns (e.g., SCP columns) may contain np.ndarray, which are not JSON serializable
+    # We need to drop them before converting the DataFrame to JSON
+    sequential_columns = [col for col in df.columns if is_sequential(df[col])]
+    df = df.drop(columns=sequential_columns)
+    _LOG.info(f"Formatting {'target' if is_target else 'context'} columns {df.columns.tolist()} to JSON")
+    # convert date format to ISO so that it's JSON serializable
+    for col in df.columns:
+        if is_datetime64_any_dtype(df[col]):
+            df[col] = df[col].dt.strftime("%Y-%m-%dT%H:%M:%S")
+    return df.apply(lambda x: row_to_json(x, is_target=is_target), axis=1)
+
+
+def fallback_serializer(obj):
+    warnings.warn(f"{type(obj)} is not JSON serializable. Converting it to str(obj) instead.")
+    return str(obj)
+
+
+def row_to_json(row: pd.Series, is_target: bool = False) -> str:
+    row_dict = row.to_dict()
+    nested_dict = {}
+    for key, value in row_dict.items():
+        tokens = key.split(TABLE_COLUMN_INFIX, maxsplit=1)
+        if len(tokens) == 1:
+            table_name = "tgt"  # assign a default table name for the target table
+            column_name = tokens[0]  # columns from the target table
+        else:  # len(tokens) == 2
+            table_name, column_name = tokens  # prefixed columns from the context tables
+        if table_name not in nested_dict:
+            nested_dict[table_name] = {}
+        nested_dict[table_name][column_name] = value
+    if is_target:
+        # assume there is at most one item in nested_dict when is_target=True
+        nested_dict = next(iter(nested_dict.values()), {})
+    # The leading space is required to avoid inconsistent tokenization results during training and generation
+    # This ensures that '{' of the context and target json are tokenized into the same thing
+    return " " + json.dumps(
+        nested_dict,
+        ensure_ascii=False,
+        separators=None,
+        indent=None,
+        default=fallback_serializer,
+    )
+
+
+def encode_df(
+    ctx_df: pd.DataFrame,
+    ctx_columns: list[str],
+    tgt_df: pd.DataFrame | None = None,
+    tgt_columns: list[str] | None = None,
+) -> pd.DataFrame:
+    assert (tgt_df is None) == (tgt_columns is None), "tgt_df and tgt_columns must be both None or both not None"
+    df = pd.DataFrame()
+    df["ctx"] = format_df(ctx_df, columns=ctx_columns, is_target=False)
+    if tgt_df is not None and tgt_columns is not None:
+        df["tgt"] = format_df(tgt_df, columns=tgt_columns, is_target=True)
+
+    # log the bounds of n_tokens in this partition
+    content = df["ctx"] + df["tgt"] if "tgt" in df.columns else df["ctx"]
+    n_chars_stats = content.apply(lambda x: len(x)).describe(percentiles=[0.5]).rename("#chars")
+    pretokenizer = ByteLevel(add_prefix_space=True, use_regex=True)
+    n_pretokens_stats = (
+        content.apply(lambda x: len(pretokenizer.pre_tokenize_str(x))).describe(percentiles=[0.5]).rename("#pretokens")
+    )
+    stats = pd.concat([n_pretokens_stats, n_chars_stats], axis=1).loc[["min", "50%", "max"], :]
+    _LOG.info(f"token statistics of this partition: \n{stats}")
+    return df
+
+
+def _encode_partition(
+    *,
+    tgt_partition_file: Path,
+    tgt_stats: dict,
+    output_path: Path,
+    ctx_partition_file: Path | None = None,
+    ctx_stats: dict | None = None,
+) -> None:
+    tgt_df = pd.read_parquet(tgt_partition_file)
+    tgt_columns = list(tgt_stats.get("columns", {}).keys())
+    if ctx_partition_file:
+        ctx_df = pd.read_parquet(ctx_partition_file)
+        ctx_columns = list(ctx_stats.get("columns", {}).keys())
+    else:
+        # create on-the-fly context
+        ctx_df = pd.DataFrame(index=range(len(tgt_df)))
+        ctx_columns = []
+    df = encode_df(
+        ctx_df=ctx_df,
+        ctx_columns=ctx_columns,
+        tgt_df=tgt_df,
+        tgt_columns=tgt_columns,
+    )
+    # shuffle and persist to disk as parquet files
+    df = df.sample(frac=1)
+    df.to_parquet(output_path / tgt_partition_file.name, engine="pyarrow", index=False)
+    _LOG.info(f"encoded partition {tgt_partition_file.name} {df.shape}")
+
+
+def encode(workspace_dir: str | Path | None = None, update_progress: ProgressCallback | None = None) -> None:
+    _LOG.info("ENCODE_LANGUAGE started")
+    t0 = time.time()
+    with ProgressCallbackWrapper(update_progress) as progress:
+        workspace_dir = ensure_workspace_dir(workspace_dir)
+        workspace = Workspace(workspace_dir)
+        reset_dir(workspace.encoded_data_val.path)
+        reset_dir(workspace.encoded_data_trn.path)
+        reset_dir(workspace.encoded_data_path)
+        has_context = workspace.ctx_data_path.exists()
+
+        tgt_stats = workspace.tgt_stats.read()
+        tgt_pqt_partitions = workspace.tgt_data.fetch_all()
+
+        if has_context:
+            ctx_stats = workspace.ctx_stats.read()
+            ctx_pqt_partitions = workspace.ctx_data.fetch_all()
+            if len(tgt_pqt_partitions) != len(ctx_pqt_partitions):
+                raise RuntimeError("partition files for tgt and ctx do not match")
+        else:
+            ctx_stats = None
+            ctx_pqt_partitions = None
+
+        for i in range(len(tgt_pqt_partitions)):
+            _encode_partition(
+                tgt_partition_file=tgt_pqt_partitions[i],
+                tgt_stats=tgt_stats,
+                output_path=workspace.encoded_data_path,
+                ctx_partition_file=ctx_pqt_partitions[i] if has_context else None,
+                ctx_stats=ctx_stats if has_context else None,
+            )
+            progress.update(completed=i, total=len(tgt_pqt_partitions) + 1)
+    _LOG.info(f"ENCODE_LANGUAGE finished in {time.time() - t0:.2f}s")

mostlyai/engine/_language/engine/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2025 MOSTLY AI
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

mostlyai/engine/_language/engine/base.py

@@ -0,0 +1,51 @@
+# 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 dataclasses import dataclass
+from abc import abstractmethod, ABC
+from collections.abc import Callable
+
+from formatron.formatter import FormatterBuilder
+
+
+@dataclass
+class EngineMetrics:
+    tokenize_time: float
+    generate_time: float
+
+
+class LanguageEngine(ABC):
+    @abstractmethod
+    def initialize_logits_processors(
+        self, formatter_builders: list[FormatterBuilder], vocab_processors: list[Callable] | None = None
+    ):
+        pass
+
+    @abstractmethod
+    def generate(
+        self, text: list[str], sampling_temperature: float, sampling_top_p: float
+    ) -> tuple[list[int], EngineMetrics]:
+        pass
+
+    @abstractmethod
+    def get_default_batch_size(self) -> int:
+        pass
+
+    @abstractmethod
+    def supports_json_enforcing(self) -> bool:
+        pass
+
+    @abstractmethod
+    def cleanup(self):
+        pass

mostlyai/engine/_language/engine/hf_engine.py

@@ -0,0 +1,136 @@
+# 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 os import PathLike
+import time
+from collections.abc import Callable
+from pathlib import Path
+
+import torch
+from formatron.integrations.transformers import create_formatter_logits_processor_list
+from peft import PeftModel
+
+from transformers import AutoTokenizer
+from mostlyai.engine._language.common import load_base_model_and_config
+from mostlyai.engine._language.tokenizer_utils import tokenize_fn
+from mostlyai.engine._language.formatron_utils import monkey_patch_formatron
+
+from mostlyai.engine._language.engine.base import EngineMetrics, LanguageEngine
+from formatron.formatter import FormatterBuilder
+
+
+class HuggingFaceEngine(LanguageEngine):
+    def __init__(
+        self, model_path: PathLike | str, device: torch.device, max_new_tokens: int, tokenizer_max_length: int
+    ):
+        self.device = device
+        self.max_new_tokens = max_new_tokens
+        self.tokenizer_max_length = tokenizer_max_length
+
+        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
+        )
+        if 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
+        else:
+            # only the LSTM model does not have an adapter
+            self._default_batch_size = 128
+
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            model_path,
+            padding_side="left",
+            truncation_side="left",
+            legacy=True,
+            # these must be False at initialization, as we manually add them later in tokenize_fn
+            add_bos_token=False,
+            add_eos_token=False,
+        )
+
+        # 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(
+            self.tokenizer.special_tokens_map
+        )
+        self._json_enforcing_possible = is_peft_adapter or is_trained_lstm_tokenizer
+
+        # apply all necessary monkey patches to the formatron library
+        if self._json_enforcing_possible:
+            monkey_patch_formatron()
+
+        self._logits_processors = None
+
+    def get_default_batch_size(self) -> int:
+        return self._default_batch_size
+
+    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 generate(
+        self, text: list[str], sampling_temperature: float, sampling_top_p: float
+    ) -> tuple[list[int], EngineMetrics]:
+        do_sample = sampling_temperature > 0.0
+
+        tokenize_kwargs = dict(
+            tokenizer=self.tokenizer,
+            return_tensors="pt",
+            add_bos_token=True,
+            add_eos_token=False,
+            padding=True,
+            truncation=True,
+            max_length=self.tokenizer_max_length,  # truncates input
+        )
+        t_tokenize = time.time()
+        inputs = tokenize_fn(text=text, **tokenize_kwargs).to(self.device)
+        tokenize_time = time.time() - t_tokenize
+
+        generate_kwargs = dict(
+            do_sample=do_sample,
+            max_new_tokens=self.max_new_tokens,
+            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
+            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:]
+        metrics = EngineMetrics(tokenize_time=tokenize_time, generate_time=generate_time)
+        return outputs.detach().cpu().tolist(), metrics
+
+    def cleanup(self):
+        pass

mostlyai/engine/_language/engine/vllm_engine.py

@@ -0,0 +1,187 @@
+# 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 contextlib
+from os import PathLike
+import time
+import typing
+from mostlyai.engine._language.engine.base import EngineMetrics, LanguageEngine
+import torch
+from formatron.config import EngineGenerationConfig
+from formatron.formatter import FormatterBuilder
+
+from peft import PeftConfig
+
+from transformers import AutoTokenizer, AutoConfig, PreTrainedTokenizerBase
+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.platforms import current_platform
+
+
+def cleanup_dist_env_and_memory():
+    """Copy from current main of vllm replace by import when possible"""
+    destroy_model_parallel()
+    destroy_distributed_environment()
+    with contextlib.suppress(AssertionError):
+        torch.distributed.destroy_process_group()
+    gc.collect()
+    if not current_platform.is_cpu():
+        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]:
+    """
+    Create list of formatter logits processors without creating engine vocabulary several times.
+    """
+    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]
+
+
+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.
+    """
+
+    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: 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
+        return scores
+
+
+class VLLMEngine(LanguageEngine):
+    def __init__(
+        self, model_path: PathLike | str, device: torch.device, max_new_tokens: int, tokenizer_max_length: int
+    ):
+        self.device = device
+        self.tokenizer_max_length = tokenizer_max_length
+        self.max_new_tokens = max_new_tokens
+
+        peft_config = PeftConfig.from_pretrained(model_path)
+        base_config = AutoConfig.from_pretrained(peft_config.base_model_name_or_path)
+
+        model_path = str(model_path)
+        self._lora_request = LoRARequest("adapter", 1, model_path)
+        config_max_model_len = _get_and_verify_max_len(
+            base_config, max_model_len=None, disable_sliding_window=None, sliding_window_len=None
+        )
+        self.llm = LLM(
+            model=peft_config.base_model_name_or_path,
+            tokenizer=model_path,
+            device=device.type,
+            max_model_len=min(config_max_model_len, self.tokenizer_max_length + max_new_tokens),
+            enable_lora=True,
+            # enforce_eager=True,  # results in big slowdown, but is needed when running pytest locally
+            swap_space=0,
+        )
+        self._base_logits_processors = [MaskInvalidIndicesLogitsProcessor(self.llm.get_tokenizer())]
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            model_path,
+            padding_side="left",
+            truncation_side="left",
+            legacy=True,
+            # these must be False at initialization, as we manually add them later in tokenize_fn
+            add_bos_token=False,
+            add_eos_token=False,
+        )
+        self._logits_processors = None
+
+    def get_default_batch_size(self) -> int:
+        return 192
+
+    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 generate(
+        self, text: list[str], sampling_temperature: float, sampling_top_p: float
+    ) -> tuple[list[int], EngineMetrics]:
+        tokenize_kwargs = dict(
+            tokenizer=self.tokenizer,
+            return_tensors=None,
+            add_bos_token=True,
+            add_eos_token=False,
+            padding=False,
+            truncation=True,
+            max_length=self.tokenizer_max_length,  # truncates input
+        )
+        t_tokenize = time.time()
+        inputs = tokenize_fn(text=text, **tokenize_kwargs)
+        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(
+                max_tokens=self.max_new_tokens,
+                temperature=sampling_temperature,
+                top_p=sampling_top_p,
+                logits_processors=self._base_logits_processors,
+            )
+        t_generate = time.time()
+        outputs = self.llm.generate(
+            prompts=None,
+            prompt_token_ids=inputs["input_ids"],
+            sampling_params=sampling_params,
+            use_tqdm=False,
+            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
+
+    def cleanup(self):
+        del self.llm
+        cleanup_dist_env_and_memory()

mostlyai/engine/_language/formatron_utils.py

@@ -0,0 +1,114 @@
+# 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 ValidationError
+from formatron.formatter import FormatterBuilder
+from typing import Literal
+from formatron.formats import json
+from pydantic import create_model
+from transformers import PreTrainedTokenizerBase
+
+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 monkey_patch_formatron():
+    # alter the Grammar of formatron's json schema
+    FORMATRON_WHITESPACE_MAX_REPETITIONS = 10
+    SPACE_NONTERMINAL = f"[ \t\n\r]{{0,{FORMATRON_WHITESPACE_MAX_REPETITIONS}}}"
+
+    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 get_formatter_builders(
+    *, seed_df: pd.DataFrame | None = None, size: int | None = None, unseeded_fields: list[str]
+) -> 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))
+    for _, seed_row in seed_df.iterrows():
+        formatter_builder = FormatterBuilder()
+        model_dict = {}
+        if not seed_row.empty:
+            model_dict |= {field_name: (Literal[seed_value], ...) for field_name, seed_value in seed_row.items()}
+        model_dict |= {field_name: (str, ...) for field_name in unseeded_fields}
+        schema = create_model("TargetModel", **model_dict, __base__=MostlyClassSchema)
+        formatter_builder.append_str(f"{formatter_builder.json(schema, capture_name=None)}")
+        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