mostlyai-engine 1.4.6__tar.gz → 1.4.7__tar.gz

{mostlyai_engine-1.4.6 → mostlyai_engine-1.4.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mostlyai-engine
-Version: 1.4.6
+Version: 1.4.7
 Summary: Synthetic Data Engine
 Project-URL: homepage, https://github.com/mostly-ai/mostlyai-engine
 Project-URL: repository, https://github.com/mostly-ai/mostlyai-engine

{mostlyai_engine-1.4.6 → mostlyai_engine-1.4.7}/mostlyai/engine/__init__.py

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

{mostlyai_engine-1.4.6 → mostlyai_engine-1.4.7}/mostlyai/engine/_language/common.py

@@ -18,7 +18,15 @@ from pathlib import Path
 
 import torch
 from peft import PeftConfig, prepare_model_for_kbit_training
-from transformers import AutoConfig, AutoModelForCausalLM, BitsAndBytesConfig, PretrainedConfig, PreTrainedModel
+from transformers import (
+    AutoConfig,
+    AutoModel,
+    AutoModelForCausalLM,
+    AutoModelForImageTextToText,
+    BitsAndBytesConfig,
+    PretrainedConfig,
+    PreTrainedModel,
+)
 
 from mostlyai.engine._language.lstm import LSTMFromScratchConfig
 
@@ -35,7 +43,7 @@ def is_bf16_supported(device: torch.device) -> bool:
 
 
 def get_attention_implementation(config: PretrainedConfig) -> str | None:
-    model_cls = AutoModelForCausalLM._model_mapping[type(config)]
+    model_cls = AutoModel._model_mapping[type(config)]
     attn_implementation = None
     if getattr(model_cls, "_supports_sdpa", False):
         attn_implementation = "sdpa"
@@ -91,13 +99,23 @@ def load_base_model_and_config(
     else:
         device_map = "auto"
 
-    model = AutoModelForCausalLM.from_pretrained(
+    if hasattr(config, "text_config") and hasattr(config, "vision_config"):
+        config.text_config.use_cache = use_cache
+        config.text_config.attn_implementation = attn_implementation
+        auto_model_cls = AutoModelForImageTextToText
+    elif hasattr(config, "use_cache"):
+        config.use_cache = use_cache
+        config.attn_implementation = attn_implementation
+        auto_model_cls = AutoModelForCausalLM
+    else:
+        raise ValueError("Unsupported model")
+
+    model = auto_model_cls.from_pretrained(
         model_id_or_path,
-        torch_dtype=torch_dtype,
-        attn_implementation=attn_implementation,
-        use_cache=use_cache,
+        config=config,
         device_map=device_map,
         quantization_config=quantization_config,
+        torch_dtype=torch_dtype,
     )
     if quantization_config:
         # convert all non-kbit layers to float32

{mostlyai_engine-1.4.6 → mostlyai_engine-1.4.7}/mostlyai/engine/_tabular/generation.py

@@ -758,15 +758,7 @@ def generate(
         _LOG.info(f"{gen_column_order=}")
         if not enable_flexible_generation:
             # check if resolved column order is the same as the one from training
-            trn_column_order = [SLEN_SIDX_SDEC_COLUMN] if is_sequential else []
-            trn_column_order += [
-                get_argn_name(
-                    argn_processor=tgt_stats["columns"][col][ARGN_PROCESSOR],
-                    argn_table=tgt_stats["columns"][col][ARGN_TABLE],
-                    argn_column=tgt_stats["columns"][col][ARGN_COLUMN],
-                )
-                for col in tgt_stats["columns"].keys()
-            ]
+            trn_column_order = get_columns_from_cardinalities(tgt_cardinalities)
             _LOG.info(f"{trn_column_order=}")
             if gen_column_order != trn_column_order:
                 raise ValueError(

{mostlyai_engine-1.4.6 → mostlyai_engine-1.4.7}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "mostlyai-engine"
-version = "1.4.6"
+version = "1.4.7"
 description = "Synthetic Data Engine"
 authors = [{ name = "MOSTLY AI", email = "dev@mostly.ai" }]
 requires-python = ">=3.10"