Bring DataDreamer up-to-date (#36)

* Initial set of updates

* Support latest torch

* Support latest torch #2

* Bump requirement versions

* Hide warning

* Minor training fixes

* Fix trainer warnings

* Update trainer fingerprints

* Update TRL

* Make Bedrock test pass

* Fix warnings

* Fix coverage

* Support FAISS in foreground for Linux

* Fix retriever test

* Fix vLLM test

* Fix VLLM tests

* Bump version

Author: AjayP13

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "DataDreamer"
-version = "0.35.0"
+version = "0.36.0"
 description = "Prompt. Generate Synthetic Data. Train & Align Models."
 license = "MIT"
 authors= [
@@ -66,7 +66,7 @@ warn_unused_ignores = true
 mypy_path = "src/_stubs"
 
 [[tool.mypy.overrides]]
-module = "click,wandb,wandb.*,click.testing,flaky,tensorflow,torch_xla,jax,datasets.features.features,datasets.iterable_dataset,datasets.fingerprint,datasets.builder,datasets.arrow_writer,datasets.splits,datasets.utils,datasets.utils.version,pyarrow.lib,huggingface_hub,huggingface_hub.utils._headers,huggingface_hub.utils._errors,dill,dill.source,transformers,bitsandbytes,sqlitedict,optimum.bettertransformer,optimum.bettertransformer.models,optimum.utils,transformers.utils.quantization_config,sortedcontainers,peft,psutil,ring,ctransformers,petals,petals.client.inference_session,hivemind.p2p.p2p_daemon_bindings.utils,huggingface_hub.utils,tqdm,ctransformers.transformers,vllm,litellm,litellm.llms.palm,litellm.exceptions,sentence_transformers,faiss,huggingface_hub.utils._validators,evaluate,transformers.trainer_callback,transformers.training_args,trl,guidance,sentence_transformers.models.Transformer,trl.trainer.utils,transformers.trainer_utils,setfit,joblib,setfit.modeling,transformers.utils.notebook,mistralai.models.chat_completion,accelerate.utils,accelerate.utils.constants,accelerate,transformers.trainer,sentence_transformers.util,Pyro5,Pyro5.server,Pyro5.api,Pyro5,datadreamer,huggingface_hub.repocard,transformers.trainer_pt_utils"
+module = "click,wandb,wandb.*,click.testing,flaky,tensorflow,torch_xla,jax,datasets.features.features,datasets.iterable_dataset,datasets.fingerprint,datasets.builder,datasets.arrow_writer,datasets.splits,datasets.utils,datasets.utils.version,pyarrow.lib,huggingface_hub,huggingface_hub.utils._headers,huggingface_hub.utils._errors,dill,dill.source,transformers,bitsandbytes,sqlitedict,optimum.bettertransformer,optimum.bettertransformer.models,optimum.utils,transformers.utils.quantization_config,sortedcontainers,peft,psutil,ring,ctransformers,petals,petals.client.inference_session,hivemind.p2p.p2p_daemon_bindings.utils,huggingface_hub.utils,tqdm,ctransformers.transformers,vllm,litellm,litellm.llms.palm,litellm.exceptions,sentence_transformers,faiss,huggingface_hub.utils._validators,evaluate,transformers.trainer_callback,transformers.training_args,trl,guidance,sentence_transformers.models.Transformer,trl.trainer.utils,transformers.trainer_utils,setfit,joblib,setfit.modeling,transformers.utils.notebook,mistralai.models.chat_completion,accelerate.utils,accelerate.utils.constants,accelerate,transformers.trainer,sentence_transformers.util,Pyro5,Pyro5.server,Pyro5.api,Pyro5,datadreamer,huggingface_hub.repocard,transformers.trainer_pt_utils,traitlets.utils.warnings,orjson,Pyro5.errors,sympy,tqdm.auto"
 ignore_missing_imports = true
 
 [tool.pyright]

src/_patches/__init__.py

@@ -0,0 +1,57 @@
+# An update in datasets 2.20.0 adding state_dict to IterableDataset seems to have
+# broken IterableDataset. This patch is a temporary fix until the issue is resolved.
+
+import contextlib
+from unittest.mock import patch
+
+from datasets.iterable_dataset import (
+    ArrowExamplesIterable,
+    ExamplesIterable,
+    TypedExamplesIterable,
+)
+
+__original_init_state_dict = TypedExamplesIterable._init_state_dict
+__original_examples__iter__ = ExamplesIterable.__iter__
+__original_arrowexamples__iter__ = ArrowExamplesIterable.__iter__
+_should_reset_state_dict = False
+
+
+def patched_examples__iter__(self):
+    global _should_reset_state_dict
+    if _should_reset_state_dict:
+        self._init_state_dict()
+    return __original_examples__iter__(self)
+
+
+def patched_arrowexamples__iter__(self):
+    global _should_reset_state_dict
+    if _should_reset_state_dict:
+        self._init_state_dict()
+    return __original_arrowexamples__iter__(self)
+
+
+ExamplesIterable.__iter__ = patched_examples__iter__
+ArrowExamplesIterable.__iter__ = patched_arrowexamples__iter__
+
+
+@contextlib.contextmanager
+def apply_datasets_reset_state_hack():
+    def patched_init_state_dict(self):
+        self._state_dict = None  # Set to None to ensure it is reset
+        return __original_init_state_dict(self)
+
+    with patch(
+        "datasets.iterable_dataset.TypedExamplesIterable._init_state_dict",
+        patched_init_state_dict,
+    ):
+        yield None
+
+
+def start_datasets_reset_state_hack():
+    global _should_reset_state_dict
+    _should_reset_state_dict = True
+
+
+def stop_datasets_reset_state_hack():
+    global _should_reset_state_dict
+    _should_reset_state_dict = False

src/_patches/datasets_reset_state_hack.py

@@ -0,0 +1,14 @@
+# SetFit is out-of-date with huggingface_hub and throws an error when trying to import
+# from it
+# like this: ImportError: cannot import name 'DatasetFilter' from 'huggingface_hub'
+
+# To fix this, we need to monkey patch huggingface_hub to prevent the import error
+
+from ..utils.import_utils import ignore_pydantic_warnings
+
+
+def apply_setfit_import_hack():
+    with ignore_pydantic_warnings():
+        import huggingface_hub
+
+        huggingface_hub.DatasetFilter = None

src/_patches/setfit_import_hack.py

@@ -15,13 +15,20 @@
 from sqlitedict import SqliteDict
 
 from . import logging as datadreamer_logging
+from ._patches.datasets_reset_state_hack import (
+    start_datasets_reset_state_hack,
+    stop_datasets_reset_state_hack,
+)
 from .logging import DATEFMT, logger
 from .utils.background_utils import get_thread_id
 from .utils.fs_utils import safe_fn
 from .utils.import_utils import ignore_pydantic_warnings, ignore_transformers_warnings
 
 with ignore_transformers_warnings():
     from optimum.utils import logging as optimum_logging
+    from ._patches.setfit_import_hack import apply_setfit_import_hack  # isort:skip
+
+    apply_setfit_import_hack()
     from setfit import logging as setfit_logging
     from transformers import logging as transformers_logging
 
@@ -517,6 +524,9 @@ def __enter__(self):  # noqa: C901
         )
         self._patch_tqdm()
 
+        # Activate datasets reset state hack
+        start_datasets_reset_state_hack()
+
         # Set initialized to True
         DataDreamer.ctx.instance = self
         DataDreamer.ctx.initialized = True
@@ -546,6 +556,7 @@ def __exit__(self, exc_type, exc_value, exc_tb):
 
         self._unpatch_loggers()
         self._unpatch_tqdm()
+        stop_datasets_reset_state_hack()
         processes_to_terminate = DataDreamer.ctx.background_processes
         DataDreamer.ctx = UserDict()
         if self.output_folder_path:

src/datadreamer.py

@@ -7,6 +7,7 @@
 from datasets.fingerprint import Hasher
 from pandas import DataFrame
 
+from .._patches.datasets_reset_state_hack import apply_datasets_reset_state_hack
 from ..datasets.utils import get_column_names
 from ..pickling import unpickle_transform
 
@@ -42,11 +43,14 @@ def _features(self) -> Features:
             return Features()
 
     def __iter__(self):
-        if self._pickled or self._pickled_inferred:  # type:ignore[attr-defined]
-            for row in iter(self.dataset):  # type:ignore[attr-defined]
-                yield unpickle_transform(row, features=self._features, batched=False)
-        else:
-            yield from iter(self.dataset)  # type:ignore[attr-defined]
+        with apply_datasets_reset_state_hack():
+            if self._pickled or self._pickled_inferred:  # type:ignore[attr-defined]
+                for row in iter(self.dataset):  # type:ignore[attr-defined]
+                    yield unpickle_transform(
+                        row, features=self._features, batched=False
+                    )
+            else:
+                yield from iter(self.dataset)  # type:ignore[attr-defined]
 
     def __getitem__(self, key: int | slice | str | Iterable[int]) -> Any:
         """Get a row or column from the dataset.
@@ -316,7 +320,8 @@ def cast_column(
         )
 
     def __iter__(self):
-        return iter(self.dataset)
+        with apply_datasets_reset_state_hack():
+            return iter(self.dataset)
 
     def __len__(self) -> int:
         return self.total_num_rows

src/datasets/datasets.py

@@ -97,7 +97,7 @@ def model(self) -> SentenceTransformer:
         model = cls(
             self.model_name,
             trust_remote_code=self.trust_remote_code,
-            device=self.device,
+            device=self.device,  # type:ignore[arg-type]
             **self.kwargs,
         )
         model[0].tokenizer = get_tokenizer(
@@ -160,7 +160,11 @@ def model_max_length(self) -> int:
 
     @cached_property
     def dims(self) -> int:
-        return self.model.get_sentence_embedding_dimension()
+        dims = self.model.get_sentence_embedding_dimension()
+        assert (
+            dims is not None
+        ), f"Failed to get the embedding dimension for {self.model_name}."
+        return dims
 
     @torch.no_grad()
     def _run_batch(
@@ -181,8 +185,8 @@ def _run_batch(
             model_input = [[cast(str, instruction), t] for t in texts]
 
         return list(
-            self.model.encode(
-                sentences=model_input,
+            self.model.encode(  # type:ignore[arg-type]
+                sentences=model_input,  # type:ignore[arg-type]
                 batch_size=len(texts),
                 show_progress_bar=False,
                 convert_to_numpy=True,

src/embedders/sentence_transformers_embedder.py

@@ -58,6 +58,7 @@ def retry_wrapper(self):
             from litellm.exceptions import (
                 APIConnectionError,
                 APIError,
+                InternalServerError,
                 RateLimitError,
                 ServiceUnavailableError,
             )
@@ -81,6 +82,14 @@ def retry_wrapper(self):
             stop=stop_any(lambda _: not self.retry_on_fail),  # type: ignore[arg-type]
             reraise=True,
         )
+        @retry(
+            retry=retry_if_exception_type(InternalServerError),
+            wait=wait_exponential(multiplier=1, min=3, max=300),
+            before_sleep=before_sleep_log(tenacity_logger, logging.INFO),
+            after=after_log(tenacity_logger, logging.INFO),
+            stop=stop_any(lambda _: not self.retry_on_fail),  # type: ignore[arg-type]
+            reraise=True,
+        )
         @retry(
             retry=retry_if_exception_type(APIError),
             wait=wait_exponential(multiplier=1, min=3, max=300),
@@ -98,7 +107,8 @@ def retry_wrapper(self):
             reraise=True,
         )
         def _retry_wrapper(func, **kwargs):
-            return func(**kwargs)
+            with ignore_litellm_warnings():
+                return func(**kwargs)
 
         _retry_wrapper.__wrapped__.__module__ = None  # type: ignore[attr-defined]
         _retry_wrapper.__wrapped__.__qualname__ = f"{self.__class__.__name__}.run"  # type: ignore[attr-defined]
@@ -126,7 +136,11 @@ def get_max_context_length(self, max_new_tokens: int) -> int:
         with ignore_litellm_warnings():
             from litellm import get_max_tokens
 
-        return get_max_tokens(model=self._model_name_prefix + self.model_name)
+        max_tokens = get_max_tokens(model=self._model_name_prefix + self.model_name)
+        assert (
+            max_tokens is not None
+        ), f"Failed to get the maximum context length for model: {self.model_name}."
+        return max_tokens
 
     @ring.lru(maxsize=5000)
     def count_tokens(self, value: str) -> int:

src/llms/_litellm.py

@@ -4,18 +4,25 @@
 from functools import cached_property, partial
 from typing import Any, Callable, Generator, Iterable
 
-import dill
 import torch
 from datasets.fingerprint import Hasher
 
 from .. import DataDreamer
 from ..logging import logger as datadreamer_logger
 from ..utils.arg_utils import AUTO, Default
-from ..utils.background_utils import RunIfTimeout, proxy_resource_in_background
+from ..utils.background_utils import (
+    RunIfTimeout,
+    dill_serializer,
+    proxy_resource_in_background,
+)
 from ..utils.device_utils import get_device_env_variables, is_cpu_device
 from ..utils.fs_utils import safe_fn
 from ..utils.hf_model_utils import get_tokenizer
-from ..utils.import_utils import ignore_transformers_warnings, import_module
+from ..utils.import_utils import (
+    ignore_tqdm,
+    ignore_transformers_warnings,
+    import_module,
+)
 from .hf_transformers import CachedTokenizer, HFTransformers
 from .llm import (
     DEFAULT_BATCH_SIZE,
@@ -62,6 +69,9 @@ def __init__(
             cache_folder_path=cache_folder_path,
             **kwargs,
         )
+        self.device = (
+            [self.device] if not isinstance(self.device, list) else self.device  # type:ignore[list-item]
+        )
         self.quantization = quantization
         if self.quantization is None and "-awq" in model_name.lower():
             self.quantization = "awq"
@@ -89,6 +99,10 @@ def _monkey_patch_init_logger(*args, **kwargs):
 
                     vllm_logging.init_logger = _monkey_patch_init_logger  # type:ignore[attr-defined]
                     logging.getLogger("vllm.engine.llm_engine").level = logging.ERROR
+                    logging.getLogger("vllm.config").level = logging.ERROR
+                    logging.getLogger(
+                        "vllm.distributed.parallel_state"
+                    ).level = logging.ERROR
 
                 # Load model
                 log_if_timeout = RunIfTimeout(
@@ -101,18 +115,19 @@ def _monkey_patch_init_logger(*args, **kwargs):
                     timeout=10.0,
                 )
                 LLM = import_module("vllm").LLM
-                self_resource.model = LLM(
-                    model=self.model_name,
-                    trust_remote_code=self.trust_remote_code,
-                    dtype=str(self.dtype).replace("torch.", "")
-                    if self.dtype is not None
-                    else "auto",
-                    quantization=self.quantization,
-                    revision=self.revision,
-                    swap_space=self.swap_space,
-                    tensor_parallel_size=tensor_parallel_size,
-                    **kwargs,
-                )
+                with ignore_tqdm():
+                    self_resource.model = LLM(
+                        model=self.model_name,
+                        trust_remote_code=self.trust_remote_code,
+                        dtype=str(self.dtype).replace("torch.", "")
+                        if self.dtype is not None
+                        else "auto",
+                        quantization=self.quantization,
+                        revision=self.revision,
+                        swap_space=self.swap_space,
+                        tensor_parallel_size=tensor_parallel_size,
+                        **kwargs,
+                    )
 
                 # Finished loading
                 log_if_timeout.stop(
@@ -124,9 +139,8 @@ def _monkey_patch_init_logger(*args, **kwargs):
                     )
                 )
 
-            def get_generated_texts_batch(self_resource, args, kwargs):
-                args = dill.loads(args)
-                kwargs = dill.loads(kwargs)
+            @dill_serializer
+            def get_generated_texts_batch(self_resource, *args, **kwargs):
                 outputs = self_resource.model.generate(*args, **kwargs)
                 generated_texts_batch = [
                     [o.text for o in batch.outputs] for batch in outputs
@@ -202,8 +216,7 @@ def _run_batch(  # noqa: C901
             **kwargs,
         )
         generated_texts_batch = self.model.proxy.get_generated_texts_batch(
-            args=dill.dumps((prompts, sampling_params)),
-            kwargs=dill.dumps({"use_tqdm": False}),
+            prompts, sampling_params, use_tqdm=False
         )
 
         # Post-process and return

src/llms/vllm.py

@@ -1 +1 @@
-torch==2.1.2,<3.0.0
+torch>=2.1.2,<3.0.0