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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mostlyai-engine
-Version: 1.0.3
+Version: 1.0.4
 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.0.3 → mostlyai_engine-1.0.4}/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.0.4"
 
 # suppress specific warning related to os.fork() in multi-threaded processes
 warnings.filterwarnings("ignore", category=DeprecationWarning, message=".*multi-threaded.*fork.*")

{mostlyai_engine-1.0.3 → mostlyai_engine-1.0.4}/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
@@ -105,8 +107,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(

{mostlyai_engine-1.0.3 → mostlyai_engine-1.0.4}/mostlyai/engine/_language/generation.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import contextlib
 import json
 import os
 
@@ -27,6 +28,7 @@ from transformers import (
     PreTrainedTokenizerBase,
 )
 
+
 from mostlyai.engine._common import (
     persist_data_part,
     FixedSizeSampleBuffer,
@@ -123,7 +125,19 @@ def generate(
 ):
     _LOG.info("GENERATE_LANGUAGE started")
     t0_ = time.time()
-    with ProgressCallbackWrapper(update_progress) as progress:
+    os.environ["VLLM_LOGGING_LEVEL"] = "WARNING"
+    os.environ["VLLM_NO_DEPRECATION_WARNING"] = "1"
+
+    @contextlib.contextmanager
+    def tqdm_disabled():
+        tqdm_disable = os.getenv("TQDM_DISABLE")
+        os.environ["TQDM_DISABLE"] = "1"
+        try:
+            yield
+        finally:
+            os.environ["TQDM_DISABLE"] = tqdm_disable if tqdm_disable is not None else ""
+
+    with ProgressCallbackWrapper(update_progress) as progress, tqdm_disabled():
         device = (
             torch.device(device)
             if device is not None

{mostlyai_engine-1.0.3 → mostlyai_engine-1.0.4}/pyproject.toml

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