implement peft lazy importing

Author: voorhs

pyproject.toml

@@ -44,14 +44,14 @@ dependencies = [
     "xxhash (>=3.5.0,<4.0.0)",
     "python-dotenv (>=1.0.1,<2.0.0)",
     "transformers[torch] (>=4.49.0,<5.0.0)",
-    "peft (>= 0.10.0, !=0.15.0, !=0.15.1, <1.0.0)",
     "aiometer (>=1.0.0,<2.0.0)",
     "aiofiles (>=24.1.0,<25.0.0)",
     "threadpoolctl (>=3.0.0,<4.0.0)",
 ]
 
 [project.optional-dependencies]
 catboost = ["catboost (>=1.2.8,<2.0.0)"]
+peft = ["peft (>= 0.10.0, !=0.15.0, !=0.15.1, <1.0.0)"]
 dspy = [
     "dspy (>=2.6.5,<3.0.0)",
 ]

src/autointent/_dump_tools/unit_dumpers.py

@@ -8,7 +8,6 @@
 import joblib
 import numpy as np
 import numpy.typing as npt
-from peft import PeftModel
 from pydantic import BaseModel
 from sklearn.base import BaseEstimator
 from transformers import (
@@ -28,6 +27,7 @@
 
 if TYPE_CHECKING:
     from catboost import CatBoostClassifier
+    from peft import PeftModel
 
 T = TypeVar("T")
 logger = logging.getLogger(__name__)
@@ -207,11 +207,11 @@ def check_isinstance(cls, obj: Any) -> bool:  # noqa: ANN401
         return isinstance(obj, BaseModel)
 
 
-class PeftModelDumper(BaseObjectDumper[PeftModel]):
+class PeftModelDumper(BaseObjectDumper["PeftModel"]):
     dir_or_file_name = "peft_models"
 
     @staticmethod
-    def dump(obj: PeftModel, path: Path, exists_ok: bool) -> None:
+    def dump(obj: "PeftModel", path: Path, exists_ok: bool) -> None:
         path.mkdir(parents=True, exist_ok=exists_ok)
         if obj._is_prompt_learning:  # noqa: SLF001
             # strategy to save prompt learning models: save prompt encoder and bert classifier separately
@@ -227,12 +227,13 @@ def dump(obj: PeftModel, path: Path, exists_ok: bool) -> None:
             merged_model.save_pretrained(lora_path)
 
     @staticmethod
-    def load(path: Path, **kwargs: Any) -> PeftModel:  # noqa: ANN401, ARG004
+    def load(path: Path, **kwargs: Any) -> "PeftModel":  # noqa: ANN401, ARG004
+        peft = require("peft", extra="peft")
         if (path / "ptuning").exists():
             # prompt learning model
             ptuning_path = path / "ptuning"
             model = AutoModelForSequenceClassification.from_pretrained(ptuning_path / "base_model")
-            return PeftModel.from_pretrained(model, ptuning_path / "peft")
+            return peft.PeftModel.from_pretrained(model, ptuning_path / "peft")
         if (path / "lora").exists():
             # merged lora model
             lora_path = path / "lora"
@@ -242,7 +243,11 @@ def load(path: Path, **kwargs: Any) -> PeftModel:  # noqa: ANN401, ARG004
 
     @classmethod
     def check_isinstance(cls, obj: Any) -> bool:  # noqa: ANN401
-        return isinstance(obj, PeftModel)
+        try:
+            peft = require("peft", extra="peft")
+            return isinstance(obj, peft.PeftModel)
+        except ImportError:
+            return False
 
 
 class HFModelDumper(BaseObjectDumper[PreTrainedModel]):

src/autointent/modules/scoring/_lora/lora.py

@@ -1,16 +1,18 @@
 """BertScorer class for transformer-based classification with LoRA."""
 
 from pathlib import Path
-from typing import Any, Literal
-
-from peft import LoraConfig, get_peft_model
+from typing import TYPE_CHECKING, Any, Literal
 
 from autointent import Context
 from autointent._callbacks import REPORTERS_NAMES
 from autointent._dump_tools import Dumper
+from autointent._utils import require
 from autointent.configs import EarlyStoppingConfig, HFModelConfig
 from autointent.modules.scoring._bert import BertScorer
 
+if TYPE_CHECKING:
+    from peft import LoraConfig
+
 
 class BERTLoRAScorer(BertScorer):
     """BERTLoRAScorer class for transformer-based classification with LoRA (Low-Rank Adaptation).
@@ -56,6 +58,8 @@ class BERTLoRAScorer(BertScorer):
 
     name = "lora"
 
+    _lora_config: "LoraConfig"
+
     def __init__(
         self,
         classification_model_config: HFModelConfig | str | dict[str, Any] | None = None,
@@ -67,6 +71,11 @@ def __init__(
         print_progress: bool = False,
         **lora_kwargs: Any,  # noqa: ANN401
     ) -> None:
+        # Lazy import peft
+        peft = require("peft", extra="peft")
+        self._LoraConfig = peft.LoraConfig
+        self._get_peft_model = peft.get_peft_model
+
         # early stopping doesnt work with lora for now https://github.com/huggingface/transformers/issues/38130
         early_stopping_config = EarlyStoppingConfig(metric=None)  # disable early stopping
 
@@ -80,7 +89,7 @@ def __init__(
             early_stopping_config=early_stopping_config,
             print_progress=print_progress,
         )
-        self._lora_config = LoraConfig(**lora_kwargs)
+        self._lora_config = self._LoraConfig(**lora_kwargs)
 
     @classmethod
     def from_context(
@@ -107,7 +116,7 @@ def from_context(
 
     def _initialize_model(self) -> Any:  # noqa: ANN401
         model = super()._initialize_model()
-        return get_peft_model(model, self._lora_config)
+        return self._get_peft_model(model, self._lora_config)
 
     def dump(self, path: str) -> None:
-        Dumper.dump(self, Path(path), exclude=[LoraConfig])
+        Dumper.dump(self, Path(path), exclude=[self._LoraConfig])

src/autointent/modules/scoring/_ptuning/ptuning.py

@@ -1,17 +1,20 @@
 """PTuningScorer class for ptuning-based classification."""
 
 from pathlib import Path
-from typing import Any, Literal
+from typing import TYPE_CHECKING, Any, Literal
 
-from peft import PromptEncoderConfig, PromptEncoderReparameterizationType, TaskType, get_peft_model
 from pydantic import PositiveInt
 
 from autointent import Context
 from autointent._callbacks import REPORTERS_NAMES
 from autointent._dump_tools import Dumper
+from autointent._utils import require
 from autointent.configs import EarlyStoppingConfig, HFModelConfig
 from autointent.modules.scoring._bert import BertScorer
 
+if TYPE_CHECKING:
+    from peft import PromptEncoderConfig
+
 
 class PTuningScorer(BertScorer):
     """PEFT P-tuning scorer.
@@ -47,6 +50,8 @@ class PTuningScorer(BertScorer):
 
     name = "ptuning"
 
+    _ptuning_config: "PromptEncoderConfig"
+
     def __init__(  # noqa: PLR0913
         self,
         classification_model_config: HFModelConfig | str | dict[str, Any] | None = None,
@@ -64,6 +69,13 @@ def __init__(  # noqa: PLR0913
         print_progress: bool = False,
         **ptuning_kwargs: Any,  # noqa: ANN401
     ) -> None:
+        # Lazy import peft
+        peft = require("peft", extra="peft")
+        self._PromptEncoderConfig = peft.PromptEncoderConfig
+        self._PromptEncoderReparameterizationType = peft.PromptEncoderReparameterizationType
+        self._TaskType = peft.TaskType
+        self._get_peft_model = peft.get_peft_model
+
         super().__init__(
             classification_model_config=classification_model_config,
             num_train_epochs=num_train_epochs,
@@ -74,9 +86,9 @@ def __init__(  # noqa: PLR0913
             early_stopping_config=early_stopping_config,
             print_progress=print_progress,
         )
-        self._ptuning_config = PromptEncoderConfig(
-            task_type=TaskType.SEQ_CLS,
-            encoder_reparameterization_type=PromptEncoderReparameterizationType(encoder_reparameterization_type),
+        self._ptuning_config = self._PromptEncoderConfig(
+            task_type=self._TaskType.SEQ_CLS,
+            encoder_reparameterization_type=self._PromptEncoderReparameterizationType(encoder_reparameterization_type),
             num_virtual_tokens=num_virtual_tokens,
             encoder_dropout=encoder_dropout,
             encoder_hidden_size=encoder_hidden_size,
@@ -139,7 +151,7 @@ def from_context(  # noqa: PLR0913
     def _initialize_model(self) -> Any:  # noqa: ANN401
         """Initialize the model with P-tuning configuration."""
         model = super()._initialize_model()
-        return get_peft_model(model, self._ptuning_config)
+        return self._get_peft_model(model, self._ptuning_config)
 
     def dump(self, path: str) -> None:
-        Dumper.dump(self, Path(path), exclude=[PromptEncoderConfig])
+        Dumper.dump(self, Path(path), exclude=[self._PromptEncoderConfig])

tests/modules/scoring/test_lora.py

@@ -8,6 +8,8 @@
 from autointent.context.data_handler import DataHandler
 from autointent.modules import BERTLoRAScorer
 
+pytest.importorskip("peft")
+
 
 def test_lora_scorer_dump_load(dataset):
     """Test that BERTLoRAScorer can be saved and loaded while preserving predictions."""

tests/modules/scoring/test_ptuning.py

@@ -8,6 +8,8 @@
 from autointent.context.data_handler import DataHandler
 from autointent.modules import PTuningScorer
 
+pytest.importorskip("peft")
+
 
 def test_ptuning_scorer_dump_load(dataset):
     """Test that PTuningScorer can be saved and loaded while preserving predictions."""