add joblib library

Author: Clara Luise Pohland

setup.py

@@ -89,7 +89,7 @@
     "mergekit": ["mergekit>=0.0.5.1"],
     "peft": ["peft>=0.8.0"],
     "quantization": ["bitsandbytes"],
-    "scikit": ["scikit-learn"],
+    "scikit": ["scikit-learn", "joblib"],
     "test": ["parameterized", "pytest-cov", "pytest-rerunfailures", "pytest-xdist", "pytest"],
     "vllm": ["vllm>=0.7.1; sys_platform != 'win32'"],  # vllm is not available on Windows
     "vlm": ["Pillow"],

tests/testing_utils.py

@@ -18,7 +18,7 @@
 from transformers import is_bitsandbytes_available, is_comet_available, is_sklearn_available, is_wandb_available
 
 from trl import BaseBinaryJudge, BasePairwiseJudge, is_diffusers_available, is_llm_blender_available
-from trl.import_utils import is_mergekit_available
+from trl.import_utils import is_joblib_available, is_mergekit_available
 
 
 # transformers.testing_utils contains a require_bitsandbytes function, but relies on pytest markers which we don't use
@@ -62,7 +62,7 @@ def require_sklearn(test_case):
     """
     Decorator marking a test that requires sklearn. Skips the test if sklearn is not available.
     """
-    return unittest.skipUnless(is_sklearn_available(), "test requires sklearn")(test_case)
+    return unittest.skipUnless(is_sklearn_available() and is_joblib_available(), "test requires sklearn")(test_case)
 
 
 def require_comet(test_case):

trl/import_utils.py

@@ -29,6 +29,7 @@
 _rich_available = _is_package_available("rich")
 _unsloth_available = _is_package_available("unsloth")
 _vllm_available = _is_package_available("vllm")
+_joblib_available = _is_package_available("joblib")
 
 
 def is_deepspeed_available() -> bool:
@@ -59,6 +60,10 @@ def is_vllm_available() -> bool:
     return _vllm_available
 
 
+def is_joblib_available() -> bool:
+    return _joblib_available
+
+
 class _LazyModule(ModuleType):
     """
     Module class that surfaces all objects but only performs associated imports when the objects are requested.

trl/trainer/bco_trainer.py

@@ -23,7 +23,6 @@
 from operator import itemgetter
 from typing import TYPE_CHECKING, Any, Callable, Literal, Optional, Union
 
-import joblib
 import numpy as np
 import pandas as pd
 import torch
@@ -56,6 +55,7 @@
 from transformers.utils import is_peft_available
 
 from ..data_utils import maybe_apply_chat_template
+from ..import_utils import is_joblib_available
 from ..models import PreTrainedModelWrapper, create_reference_model
 from .bco_config import BCOConfig
 from .utils import (
@@ -80,6 +80,9 @@
 if is_sklearn_available():
     from sklearn.linear_model import LogisticRegression
 
+if is_joblib_available():
+    import joblib
+
 if is_deepspeed_available():
     import deepspeed
 
@@ -350,9 +353,9 @@ def __init__(
         embedding_func: Optional[Callable] = None,
         embedding_tokenizer: Optional[PreTrainedTokenizerBase] = None,
     ):
-        if not is_sklearn_available():
+        if embedding_func != None and not (is_sklearn_available() and is_joblib_available()):
             raise ImportError(
-                "BCOTrainer requires the scikit-learn library. Please install it with `pip install scikit-learn`."
+                "BCOTrainer with UDM requires the scikit-learn and joblib libraries. Please install it with `pip install scikit-learn joblib`."
             )
 
         if type(args) is TrainingArguments: