mostlyai-engine 1.0.3__tar.gz → 1.1.0__tar.gz

{mostlyai_engine-1.0.3 → mostlyai_engine-1.1.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mostlyai-engine
-Version: 1.0.3
+Version: 1.1.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
@@ -37,7 +37,7 @@ 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.21,>=0.20.1
+Requires-Dist: tokenizers<0.22,>=0.21.0
 Requires-Dist: torch>=2.5.1; sys_platform == 'darwin'
 Requires-Dist: transformers<5,>=4.45.2
 Provides-Extra: cpu
@@ -46,7 +46,7 @@ Requires-Dist: torch>=2.5.1; (sys_platform != 'linux') and extra == 'cpu'
 Provides-Extra: gpu
 Requires-Dist: bitsandbytes<0.44,>=0.43.3; extra == 'gpu'
 Requires-Dist: torch>=2.5.1; extra == 'gpu'
-Requires-Dist: vllm<0.7,>=0.6.4; extra == 'gpu'
+Requires-Dist: vllm<0.8,>=0.7.2; extra == 'gpu'
 Description-Content-Type: text/markdown
 
 # Synthetic Data Engine 💎

{mostlyai_engine-1.0.3 → mostlyai_engine-1.1.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.0.3"
+__version__ = "1.1.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.0/mostlyai/engine/_encoding_types/language/categorical.py

@@ -0,0 +1,75 @@
+# 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.
+
+"""
+Categorical encoding for language models.
+"""
+
+import numpy as np
+import pandas as pd
+
+from mostlyai.engine._common import safe_convert_string, STRING
+
+CATEGORICAL_UNKNOWN_TOKEN = "_RARE_"
+
+
+def analyze_language_categorical(values: pd.Series, root_keys: pd.Series, _: pd.Series | None = None) -> dict:
+    values = safe_convert_string(values)
+    # count distinct root_keys per categorical value for rare-category protection
+    df = pd.concat([root_keys, values], axis=1)
+    cnt_values = df.groupby(values.name)[root_keys.name].nunique().to_dict()
+    stats = {"has_nan": sum(values.isna()) > 0, "cnt_values": cnt_values}
+    return stats
+
+
+def analyze_reduce_language_categorical(stats_list: list[dict], value_protection: bool = True) -> dict:
+    # sum up all counts for each categorical value
+    cnt_values: dict[str, int] = {}
+    for item in stats_list:
+        for value, count in item["cnt_values"].items():
+            cnt_values[value] = cnt_values.get(value, 0) + count
+    # create alphabetically sorted list of non-rare categories
+    known_categories = [k for k in sorted(cnt_values.keys())]
+    if value_protection:
+        # stochastic threshold for rare categories
+        rare_min = 5 + int(3 * np.random.uniform())
+    else:
+        rare_min = 0
+    categories = [k for k in known_categories if cnt_values[k] >= rare_min]
+    no_of_rare_categories = len(known_categories) - len(categories)
+    # add None to categories, if any are present
+    if any([j["has_nan"] for j in stats_list]):
+        categories = [None] + categories
+    # add special token for UNKNOWN categories at first position
+    if no_of_rare_categories > 0:
+        categories = [CATEGORICAL_UNKNOWN_TOKEN] + categories
+    stats = {"no_of_rare_categories": no_of_rare_categories, "categories": categories}
+    return stats
+
+
+def encode_language_categorical(values: pd.Series, stats: dict) -> pd.Series:
+    values = safe_convert_string(values)
+    values = values.copy()
+    known_categories = stats["categories"]
+    mask = ~values.isin(known_categories)
+    if None in known_categories:
+        mask &= ~pd.isna(values)
+    values[mask] = CATEGORICAL_UNKNOWN_TOKEN
+    return values
+
+
+def decode_language_categorical(x: pd.Series, col_stats: dict[str, str]) -> pd.Series:
+    x = x.astype(STRING)
+    allowed_categories = col_stats.get("categories", [])
+    return x.where(x.isin(allowed_categories), other=None)

mostlyai_engine-1.1.0/mostlyai/engine/_encoding_types/language/datetime.py

@@ -0,0 +1,143 @@
+# 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 calendar
+
+import numpy as np
+import pandas as pd
+
+from mostlyai.engine._common import safe_convert_datetime
+
+
+def analyze_language_datetime(values: pd.Series, root_keys: pd.Series, _: pd.Series | None = None) -> dict:
+    values = safe_convert_datetime(values)
+    df = pd.concat([root_keys, values], axis=1)
+    # determine lowest/highest values by root ID, and return Top 10
+    min_dates = df.groupby(root_keys.name)[values.name].min().dropna()
+    min11 = min_dates.sort_values(ascending=True).head(11).astype(str).tolist()
+    max_dates = df.groupby(root_keys.name)[values.name].max().dropna()
+    max11 = max_dates.sort_values(ascending=False).head(11).astype(str).tolist()
+    # determine if there are any NaN values
+    has_nan = bool(values.isna().any())
+    # return stats
+    stats = {
+        "has_nan": has_nan,
+        "min11": min11,
+        "max11": max11,
+    }
+    return stats
+
+
+def analyze_reduce_language_datetime(stats_list: list[dict], value_protection: bool = True) -> dict:
+    # check if there are missing values
+    has_nan = any([j["has_nan"] for j in stats_list])
+    # determine min / max 5 values to map too low / too high values to
+    min11 = sorted([v for min11 in [j["min11"] for j in stats_list] for v in min11], reverse=False)[:11]
+    max11 = sorted([v for max11 in [j["max11"] for j in stats_list] for v in max11], reverse=True)[:11]
+    if value_protection:
+        # extreme value protection - discard lowest/highest 5 values
+        if len(min11) < 11 or len(max11) < 11:
+            # less than 11 subjects with non-NULL values; we need to protect all
+            min5 = []
+            max5 = []
+        else:
+            min5 = [str(v) for v in min11[5:10]]  # drop 1 to 5th lowest; keep 6th to 10th lowest
+            max5 = [str(v) for v in max11[5:10]]  # drop 1 to 5th highest; keep 6th to 10th highest
+    else:
+        min5 = min11[0:4]
+        max5 = max11[0:4]
+    stats = {
+        "has_nan": has_nan,
+        "min5": min5,
+        "max5": max5,
+    }
+    return stats
+
+
+def encode_language_datetime(values: pd.Series, stats: dict, _: pd.Series | None = None) -> pd.Series:
+    # convert
+    values = safe_convert_datetime(values)
+    values = values.copy()
+    # reset index, as `values.mask` can throw errors for misaligned indices
+    values.reset_index(drop=True, inplace=True)
+    # replace extreme values with randomly sampled 5-th to 10-th largest/smallest values
+    min5 = stats["min5"] if len(stats["min5"]) > 0 else [0]
+    max5 = stats["max5"] if len(stats["max5"]) > 0 else [0]
+    min5 = pd.Series(min5, dtype=values.dtype)
+    max5 = pd.Series(max5, dtype=values.dtype)
+    values.mask(
+        values < min5[0],
+        min5.sample(n=len(values), replace=True, ignore_index=True),
+        inplace=True,
+    )
+    values.mask(
+        values > max5[0],
+        max5.sample(n=len(values), replace=True, ignore_index=True),
+        inplace=True,
+    )
+    return values
+
+
+def _clip_datetime(x: pd.Series, min5: list, max5: list) -> pd.Series:
+    x_dt = pd.to_datetime(x, errors="coerce")
+    min_arr = pd.to_datetime(min5).to_numpy(dtype="datetime64[ns]")
+    max_arr = pd.to_datetime(max5).to_numpy(dtype="datetime64[ns]")
+    n = len(x_dt)
+    random_mins = np.random.choice(min_arr, size=n)
+    random_maxs = np.random.choice(max_arr, size=n)
+    clipped = np.minimum(np.maximum(x_dt.to_numpy(dtype="datetime64[ns]"), random_mins), random_maxs)
+    return pd.Series(clipped, index=x.index)
+
+
+def decode_language_datetime(x: pd.Series, col_stats: dict[str, str]) -> pd.Series:
+    x = x.where(~x.isin(["", "_INVALID_"]), np.nan)
+
+    valid_mask = (
+        x.str.len().ge(10)
+        & x.str.slice(0, 4).str.isdigit()
+        & x.str.slice(5, 7).str.isdigit()
+        & x.str.slice(8, 10).str.isdigit()
+    )
+    if valid_mask.sum() > 0:  # expected "YYYY-MM-DD" prefix
+        # handle the date portion, ensuring validity
+        years = x[valid_mask].str.slice(0, 4).astype(int)
+        months = x[valid_mask].str.slice(5, 7).astype(int)
+        days = x[valid_mask].str.slice(8, 10).astype(int)
+
+        # clamp days according to maximum possible day of the month of a given year
+        last_days = np.array([calendar.monthrange(y, m)[1] for y, m in zip(years, months)])
+        clamped_days = np.minimum(days, last_days)
+
+        # rebuild the date portion
+        new_date = (
+            years.astype(str).str.zfill(4)
+            + "-"
+            + months.astype(str).str.zfill(2)
+            + "-"
+            + pd.Series(clamped_days, index=years.index).astype(str).str.zfill(2)
+        )
+
+        # handle the time portion, ensuring validity
+        remainder = x[valid_mask].str.slice(10)
+
+        time_regex = r"^[ T]?(\d{2}:\d{2}:\d{2}(?:\.\d+)?)"
+        valid_time = remainder.str.extract(time_regex, expand=False)
+        valid_time = valid_time.fillna("00:00:00")
+        valid_time = " " + valid_time
+
+        new_date = new_date + valid_time
+        x.loc[valid_mask] = new_date
+
+    x = pd.to_datetime(x, errors="coerce")
+    x = _clip_datetime(x, col_stats["min5"], col_stats["max5"])
+    return x.astype("datetime64[ns]")

mostlyai_engine-1.1.0/mostlyai/engine/_encoding_types/language/numeric.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.
+import numpy as np
+import pandas as pd
+
+from mostlyai.engine._common import safe_convert_numeric
+from mostlyai.engine._encoding_types.tabular.numeric import _type_safe_numeric_series
+from mostlyai.engine.domain import ModelEncodingType
+
+
+def analyze_language_numeric(values: pd.Series, root_keys: pd.Series, _: pd.Series | None = None) -> dict:
+    values = safe_convert_numeric(values)
+
+    # determine lowest/highest values by root ID, and return top 11
+    df = pd.concat([root_keys, values], axis=1)
+    min_values = df.groupby(root_keys.name)[values.name].min().dropna()
+    min11 = min_values.sort_values(ascending=True).head(11).tolist()
+    max_values = df.groupby(root_keys.name)[values.name].max().dropna()
+    max11 = max_values.sort_values(ascending=False).head(11).tolist()
+
+    # determine if there are any NaN values
+    has_nan = bool(values.isna().any())
+
+    # determine max scale
+    def count_scale(num: float) -> int:
+        # represent number as fixed point string, remove trailing zeros and decimal point
+        num = format(num, "f").rstrip("0").rstrip(".")
+        if "." in num:
+            # in case of decimal, return number of digits after decimal point
+            return len(num.split(".")[1])
+        # in case of integer, return 0
+        return 0
+
+    max_scale = int(values.apply(count_scale).max())
+
+    stats = {
+        "has_nan": has_nan,
+        "max_scale": max_scale,
+        "min11": min11,
+        "max11": max11,
+    }
+    return stats
+
+
+def analyze_reduce_language_numeric(stats_list: list[dict], value_protection: bool = True) -> dict:
+    # check for occurrence of NaN values
+    has_nan = any([j["has_nan"] for j in stats_list])
+
+    # determine max scale
+    max_scale = max([j["max_scale"] for j in stats_list])
+
+    # determine min / max 5 values to map too low / too high values to
+    min11 = sorted([v for min11 in [j["min11"] for j in stats_list] for v in min11], reverse=False)[:11]
+    max11 = sorted([v for max11 in [j["max11"] for j in stats_list] for v in max11], reverse=True)[:11]
+    if value_protection:
+        # extreme value protection - discard lowest/highest 5 values
+        if len(min11) < 11 or len(max11) < 11:
+            # less than 11 subjects with non-NULL values; we need to protect all
+            min5 = []
+            max5 = []
+        else:
+            min5 = min11[5:10]  # drop 1 to 5th lowest; keep 6th to 10th lowest
+            max5 = max11[5:10]  # drop 1 to 5th highest; keep 6th to 10th highest
+    else:
+        min5 = min11[0:5]
+        max5 = max11[0:5]
+
+    stats = {
+        "encoding_type": ModelEncodingType.language_numeric.value,
+        "has_nan": has_nan,
+        "max_scale": max_scale,
+        "min5": min5,
+        "max5": max5,
+    }
+
+    return stats
+
+
+def encode_language_numeric(values: pd.Series, stats: dict, _: pd.Series | None = None) -> pd.DataFrame:
+    values = safe_convert_numeric(values)
+    # try to convert to int, if possible
+    dtype = "Int64" if stats["max_scale"] == 0 else "Float64"
+    if dtype == "Int64":
+        values = values.round()
+    try:
+        values = values.astype(dtype)
+    except TypeError:
+        if dtype == "Int64":  # if couldn't safely convert to int, stick to float
+            dtype = "Float64"
+            values = values.astype(dtype)
+    # reset index, as `values.mask` can throw errors for misaligned indices
+    values.reset_index(drop=True, inplace=True)
+    # replace extreme values with randomly sampled 5-th to 10-th largest/smallest values
+    min5 = _type_safe_numeric_series(stats["min5"] or [0], dtype)
+    max5 = _type_safe_numeric_series(stats["max5"] or [0], dtype)
+    values.mask(
+        values < min5[0],
+        min5.sample(n=len(values), replace=True, ignore_index=True),
+        inplace=True,
+    )
+    values.mask(
+        values > max5[0],
+        max5.sample(n=len(values), replace=True, ignore_index=True),
+        inplace=True,
+    )
+    return values
+
+
+def _clip_numeric(x: pd.Series, min5: list, max5: list) -> pd.Series:
+    x_numeric = pd.to_numeric(x, errors="coerce")
+    min_arr = np.array(min5, dtype=x_numeric.dtype)
+    max_arr = np.array(max5, dtype=x_numeric.dtype)
+    n = len(x_numeric)
+    random_mins = np.random.choice(min_arr, size=n)
+    random_maxs = np.random.choice(max_arr, size=n)
+    clipped = np.minimum(np.maximum(x_numeric.to_numpy(), random_mins), random_maxs)
+    return pd.Series(clipped, index=x.index)
+
+
+def decode_language_numeric(x: pd.Series, col_stats: dict[str, str]) -> pd.Series:
+    x = pd.to_numeric(x, errors="coerce")
+    x = x.round(col_stats["max_scale"])
+    x = _clip_numeric(x, col_stats["min5"], col_stats["max5"])
+    dtype = "Int64" if col_stats["max_scale"] == 0 else float
+    return x.astype(dtype)

{mostlyai_engine-1.0.3 → mostlyai_engine-1.1.0}/mostlyai/engine/_encoding_types/language/text.py

@@ -14,7 +14,7 @@
 
 import pandas as pd
 
-from mostlyai.engine._common import safe_convert_string
+from mostlyai.engine._common import safe_convert_string, STRING
 
 
 def analyze_text(values: pd.Series, root_keys: pd.Series, _: pd.Series | None = None) -> dict:
@@ -39,3 +39,7 @@ def analyze_reduce_text(stats_list: list[dict], _: bool = True) -> dict:
         "nchar_max": nchar_max,
     }
     return stats
+
+
+def decode_text(x: pd.Series, col_stats: dict[str, str]) -> pd.Series:
+    return x.astype(STRING)

{mostlyai_engine-1.0.3 → mostlyai_engine-1.1.0}/mostlyai/engine/_encoding_types/tabular/numeric.py

@@ -165,7 +165,7 @@ def analyze_numeric(
         # do not count values, if there are too many
         cnt_values = None
 
-    # determine lowest/highest values by root ID, and return Top 10
+    # determine lowest/highest values by root ID, and return top 11
     df = pd.concat([root_keys, values], axis=1)
     min_values = df.groupby(root_keys.name)[values.name].min().dropna()
     min11 = min_values.sort_values(ascending=True).head(11).astype("float").tolist()

{mostlyai_engine-1.0.3 → mostlyai_engine-1.1.0}/mostlyai/engine/_language/encoding.py

@@ -24,17 +24,36 @@ 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
+from mostlyai.engine._encoding_types.language.categorical import encode_language_categorical
+from mostlyai.engine._encoding_types.language.numeric import encode_language_numeric
+from mostlyai.engine._encoding_types.language.datetime import encode_language_datetime
 
 _LOG = logging.getLogger(__name__)
 
 
-def format_df(df: pd.DataFrame, columns: list[str], is_target: bool = False) -> pd.DataFrame:
-    df = df[columns].copy()
+def apply_encoding_types(df: pd.DataFrame, stats: dict) -> pd.DataFrame:
+    for col, col_stats in stats["columns"].items():
+        if col_stats["encoding_type"] == "LANGUAGE_CATEGORICAL":
+            df[col] = encode_language_categorical(df[col], col_stats)
+        elif col_stats["encoding_type"] == "LANGUAGE_NUMERIC":
+            df[col] = encode_language_numeric(df[col], col_stats)
+        elif col_stats["encoding_type"] == "LANGUAGE_DATETIME":
+            df[col] = encode_language_datetime(df[col], col_stats)
+    return df
+
+
+def drop_sequential_columns(df: pd.DataFrame) -> pd.DataFrame:
     # 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")
+    return df
+
+
+def format_df(df: pd.DataFrame, stats: dict, is_target: bool = False) -> pd.DataFrame:
+    columns = list(stats["columns"].keys())
+    df = df[columns].copy()
+    _LOG.info(f"Formatting {'target' if is_target else 'context'} columns {columns} 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]):
@@ -76,15 +95,21 @@ def row_to_json(row: pd.Series, is_target: bool = False) -> str:
 
 def encode_df(
     ctx_df: pd.DataFrame,
-    ctx_columns: list[str],
+    ctx_stats: dict | None = None,
     tgt_df: pd.DataFrame | None = None,
-    tgt_columns: list[str] | None = None,
+    tgt_stats: dict | 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"
+    assert (tgt_df is None) == (tgt_stats is None), "tgt_df and tgt_stats must be both None or both not None"
+    if ctx_stats is None:
+        ctx_stats = {"columns": {}}
     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)
+    ctx_df = drop_sequential_columns(ctx_df)
+    ctx_df = apply_encoding_types(ctx_df, stats=ctx_stats)
+    df["ctx"] = format_df(ctx_df, stats=ctx_stats, is_target=False)
+    if tgt_df is not None and tgt_stats is not None:
+        tgt_df = drop_sequential_columns(tgt_df)
+        tgt_df = apply_encoding_types(tgt_df, stats=tgt_stats)
+        df["tgt"] = format_df(tgt_df, stats=tgt_stats, 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"]
@@ -107,19 +132,16 @@ def _encode_partition(
     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,
+        ctx_stats=ctx_stats,
         tgt_df=tgt_df,
-        tgt_columns=tgt_columns,
+        tgt_stats=tgt_stats,
     )
     # shuffle and persist to disk as parquet files
     df = df.sample(frac=1)

{mostlyai_engine-1.0.3 → mostlyai_engine-1.1.0}/mostlyai/engine/_language/engine/hf_engine.py

@@ -23,8 +23,8 @@ 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.tokenizer_utils import tokenize_fn
 
 from mostlyai.engine._language.engine.base import EngineMetrics, LanguageEngine
 from formatron.formatter import FormatterBuilder
@@ -66,11 +66,8 @@ class HuggingFaceEngine(LanguageEngine):
             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:
+        if self.supports_json_enforcing():
             monkey_patch_formatron()
-
         self._logits_processors = None
 
     def get_default_batch_size(self) -> int:

{mostlyai_engine-1.0.3 → mostlyai_engine-1.1.0}/mostlyai/engine/_language/engine/vllm_engine.py

@@ -16,6 +16,8 @@ import contextlib
 from os import PathLike
 import time
 import typing
+
+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
@@ -24,6 +26,7 @@ from formatron.formatter import FormatterBuilder
 from peft import PeftConfig
 
 from transformers import AutoTokenizer, AutoConfig, PreTrainedTokenizerBase
+from mostlyai.engine._language.formatron_utils import monkey_patch_formatron
 from vllm import LLM, SamplingParams
 from vllm.lora.request import LoRARequest
 from vllm.config import _get_and_verify_max_len
@@ -105,8 +108,10 @@ class VLLMEngine(LanguageEngine):
             device=device.type,
             max_model_len=min(config_max_model_len, self.tokenizer_max_length + max_new_tokens),
             enable_lora=True,
+            dtype=torch.bfloat16 if is_bf16_supported(device) else torch.float16,
             # enforce_eager=True,  # results in big slowdown, but is needed when running pytest locally
             swap_space=0,
+            disable_log_stats=True,
         )
         self._base_logits_processors = [MaskInvalidIndicesLogitsProcessor(self.llm.get_tokenizer())]
         self.tokenizer = AutoTokenizer.from_pretrained(
@@ -119,6 +124,7 @@ class VLLMEngine(LanguageEngine):
             add_eos_token=False,
         )
         self._logits_processors = None
+        monkey_patch_formatron()
 
     def get_default_batch_size(self) -> int:
         return 192