implement lazy importing transformers

voorhs · voorhs · commit bcdf0f2321a6 · 2025-10-23T19:47:17.000+03:00
diff --git a/pyproject.toml b/pyproject.toml
@@ -52,6 +52,7 @@ dependencies = [
 [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)"]
+transformers = ["transformers[torch] (>=4.49.0,<5.0.0)"]
 dspy = [
     "dspy (>=2.6.5,<3.0.0)",
 ]
diff --git a/src/autointent/_dump_tools/unit_dumpers.py b/src/autointent/_dump_tools/unit_dumpers.py
@@ -10,13 +10,6 @@
 import numpy.typing as npt
 from pydantic import BaseModel
 from sklearn.base import BaseEstimator
-from transformers import (
-    AutoModelForSequenceClassification,
-    AutoTokenizer,
-    PreTrainedModel,
-    PreTrainedTokenizer,
-    PreTrainedTokenizerFast,
-)
 
 from autointent import Embedder, Ranker, VectorIndex
 from autointent._utils import require
@@ -28,6 +21,7 @@
 if TYPE_CHECKING:
     from catboost import CatBoostClassifier
     from peft import PeftModel
+    from transformers import PreTrainedModel, PreTrainedTokenizer, PreTrainedTokenizerFast
 
 T = TypeVar("T")
 logger = logging.getLogger(__name__)
@@ -223,21 +217,22 @@ def dump(obj: "PeftModel", path: Path, exists_ok: bool) -> None:
             # strategy to save lora models: merge adapters and save as usual hugging face model
             lora_path = path / "lora"
             lora_path.mkdir(parents=True, exist_ok=exists_ok)
-            merged_model: PreTrainedModel = obj.merge_and_unload()
+            merged_model: "PreTrainedModel" = obj.merge_and_unload()
             merged_model.save_pretrained(lora_path)
 
     @staticmethod
     def load(path: Path, **kwargs: Any) -> "PeftModel":  # noqa: ANN401, ARG004
         peft = require("peft", extra="peft")
+        transformers = require("transformers", extra="transformers")
         if (path / "ptuning").exists():
             # prompt learning model
             ptuning_path = path / "ptuning"
-            model = AutoModelForSequenceClassification.from_pretrained(ptuning_path / "base_model")
+            model = transformers.AutoModelForSequenceClassification.from_pretrained(ptuning_path / "base_model")
             return peft.PeftModel.from_pretrained(model, ptuning_path / "peft")
         if (path / "lora").exists():
             # merged lora model
             lora_path = path / "lora"
-            return AutoModelForSequenceClassification.from_pretrained(lora_path)  # type: ignore[no-any-return]
+            return transformers.AutoModelForSequenceClassification.from_pretrained(lora_path)  # type: ignore[no-any-return]
         msg = f"Invalid PeftModel directory structure at {path}. Expected 'ptuning' or 'lora' subdirectory."
         raise ValueError(msg)
 
@@ -250,38 +245,48 @@ def check_isinstance(cls, obj: Any) -> bool:  # noqa: ANN401
             return False
 
 
-class HFModelDumper(BaseObjectDumper[PreTrainedModel]):
+class HFModelDumper(BaseObjectDumper["PreTrainedModel"]):
     dir_or_file_name = "hf_models"
 
     @staticmethod
-    def dump(obj: PreTrainedModel, path: Path, exists_ok: bool) -> None:
+    def dump(obj: "PreTrainedModel", path: Path, exists_ok: bool) -> None:
         path.mkdir(parents=True, exist_ok=exists_ok)
         obj.save_pretrained(path)
 
     @staticmethod
-    def load(path: Path, **kwargs: Any) -> PreTrainedModel:  # noqa: ANN401, ARG004
-        return AutoModelForSequenceClassification.from_pretrained(path)  # type: ignore[no-any-return]
+    def load(path: Path, **kwargs: Any) -> "PreTrainedModel":  # noqa: ANN401, ARG004
+        transformers = require("transformers", extra="transformers")
+        return transformers.AutoModelForSequenceClassification.from_pretrained(path)  # type: ignore[no-any-return]
 
     @classmethod
     def check_isinstance(cls, obj: Any) -> bool:  # noqa: ANN401
-        return isinstance(obj, PreTrainedModel)
+        try:
+            transformers = require("transformers", extra="transformers")
+            return isinstance(obj, transformers.PreTrainedModel)
+        except ImportError:
+            return False
 
 
-class HFTokenizerDumper(BaseObjectDumper[PreTrainedTokenizer | PreTrainedTokenizerFast]):
+class HFTokenizerDumper(BaseObjectDumper["PreTrainedTokenizer | PreTrainedTokenizerFast"]):
     dir_or_file_name = "hf_tokenizers"
 
     @staticmethod
-    def dump(obj: PreTrainedTokenizer | PreTrainedTokenizerFast, path: Path, exists_ok: bool) -> None:
+    def dump(obj: "PreTrainedTokenizer | PreTrainedTokenizerFast", path: Path, exists_ok: bool) -> None:
         path.mkdir(parents=True, exist_ok=exists_ok)
         obj.save_pretrained(path)
 
     @staticmethod
-    def load(path: Path, **kwargs: Any) -> PreTrainedTokenizer | PreTrainedTokenizerFast:  # noqa: ANN401, ARG004
-        return AutoTokenizer.from_pretrained(path)  # type: ignore[no-any-return,no-untyped-call]
+    def load(path: Path, **kwargs: Any) -> "PreTrainedTokenizer | PreTrainedTokenizerFast":  # noqa: ANN401, ARG004
+        transformers = require("transformers", extra="transformers")
+        return transformers.AutoTokenizer.from_pretrained(path)  # type: ignore[no-any-return,no-untyped-call]
 
     @classmethod
     def check_isinstance(cls, obj: Any) -> bool:  # noqa: ANN401
-        return isinstance(obj, PreTrainedTokenizer | PreTrainedTokenizerFast)
+        try:
+            transformers = require("transformers", extra="transformers")
+            return isinstance(obj, transformers.PreTrainedTokenizer | transformers.PreTrainedTokenizerFast)
+        except ImportError:
+            return False
 
 
 class TorchModelDumper(BaseObjectDumper[BaseTorchModule]):
diff --git a/src/autointent/_wrappers/embedder/sentence_transformers.py b/src/autointent/_wrappers/embedder/sentence_transformers.py
@@ -2,7 +2,7 @@
 import tempfile
 from functools import lru_cache
 from pathlib import Path
-from typing import Literal, cast, overload
+from typing import TYPE_CHECKING, Literal, cast, overload
 from uuid import uuid4
 
 import huggingface_hub
@@ -14,16 +14,19 @@
 from sentence_transformers.losses import BatchAllTripletLoss
 from sentence_transformers.training_args import BatchSamplers
 from sklearn.model_selection import train_test_split
-from transformers import EarlyStoppingCallback, TrainerCallback
 
 from autointent._hash import Hasher
+from autointent._utils import require
 from autointent.configs import EmbedderFineTuningConfig, TaskTypeEnum
 from autointent.configs._embedder import SentenceTransformerEmbeddingConfig
 from autointent.custom_types import ListOfLabels
 
 from .base import BaseEmbeddingBackend
 from .utils import get_embeddings_path
 
+if TYPE_CHECKING:
+    from transformers import TrainerCallback
+
 logger = logging.getLogger(__name__)
 
 
@@ -234,6 +237,9 @@ def train(self, utterances: list[str], labels: ListOfLabels, config: EmbedderFin
 
         loss = BatchAllTripletLoss(model=model, margin=config.margin)
         with tempfile.TemporaryDirectory() as tmp_dir:
+            # Lazy import transformers (only needed for fine-tuning)
+            transformers = require("transformers", extra="transformers")
+
             args = SentenceTransformerTrainingArguments(
                 save_strategy="epoch",
                 save_total_limit=1,
@@ -251,8 +257,8 @@ def train(self, utterances: list[str], labels: ListOfLabels, config: EmbedderFin
                 eval_strategy="epoch",
                 greater_is_better=False,
             )
-            callbacks: list[TrainerCallback] = [
-                EarlyStoppingCallback(
+            callbacks: list["TrainerCallback"] = [
+                transformers.EarlyStoppingCallback(
                     early_stopping_patience=config.early_stopping_patience,
                     early_stopping_threshold=config.early_stopping_threshold,
                 )
diff --git a/src/autointent/context/data_handler/_stratification.py b/src/autointent/context/data_handler/_stratification.py
@@ -5,6 +5,7 @@
 """
 
 import logging
+import random
 from collections.abc import Sequence
 
 import numpy as np
@@ -13,7 +14,6 @@
 from numpy import typing as npt
 from sklearn.model_selection import train_test_split
 from skmultilearn.model_selection import IterativeStratification
-from transformers import set_seed
 
 from autointent import Dataset
 from autointent.custom_types import LabelType
@@ -156,7 +156,8 @@ def _split_multilabel(self, dataset: HFDataset, test_size: float) -> Sequence[np
             A sequence containing indices for train and test splits.
         """
         if self.random_seed is not None:
-            set_seed(self.random_seed)  # workaround for buggy nature of IterativeStratification from skmultilearn
+            # Set all seeds for reproducibility (workaround for buggy nature of IterativeStratification from skmultilearn)
+            random.seed(self.random_seed)
         splitter = IterativeStratification(
             n_splits=2,
             order=2,
diff --git a/src/autointent/modules/scoring/_bert.py b/src/autointent/modules/scoring/_bert.py
@@ -2,33 +2,25 @@
 
 import tempfile
 from collections.abc import Callable
-from typing import Any, Literal
+from typing import TYPE_CHECKING, Any, Literal
 
 import numpy as np
 import numpy.typing as npt
 import torch
 from datasets import Dataset, DatasetDict
 from sklearn.model_selection import train_test_split
-from transformers import (
-    AutoModelForSequenceClassification,
-    AutoTokenizer,
-    DataCollatorWithPadding,
-    EarlyStoppingCallback,
-    EvalPrediction,
-    PrinterCallback,
-    ProgressCallback,
-    Trainer,
-    TrainingArguments,
-)
-from transformers.trainer_callback import TrainerCallback
 
 from autointent import Context
 from autointent._callbacks import REPORTERS_NAMES
+from autointent._utils import require
 from autointent.configs import EarlyStoppingConfig, HFModelConfig
 from autointent.custom_types import ListOfLabels
 from autointent.metrics import SCORING_METRICS_MULTICLASS, SCORING_METRICS_MULTILABEL
 from autointent.modules.base import BaseScorer
 
+if TYPE_CHECKING:
+    from transformers import EvalPrediction, TrainerCallback
+
 
 class BertScorer(BaseScorer):
     """Scoring module for transformer-based classification using BERT models.
@@ -90,6 +82,17 @@ def __init__(
         early_stopping_config: EarlyStoppingConfig | dict[str, Any] | None = None,
         print_progress: bool = False,
     ) -> None:
+        # Lazy import transformers
+        transformers = require("transformers", extra="transformers")
+        self._AutoModelForSequenceClassification = transformers.AutoModelForSequenceClassification
+        self._AutoTokenizer = transformers.AutoTokenizer
+        self._DataCollatorWithPadding = transformers.DataCollatorWithPadding
+        self._EarlyStoppingCallback = transformers.EarlyStoppingCallback
+        self._PrinterCallback = transformers.PrinterCallback
+        self._ProgressCallback = transformers.ProgressCallback
+        self._Trainer = transformers.Trainer
+        self._TrainingArguments = transformers.TrainingArguments
+
         self.classification_model_config = HFModelConfig.from_search_config(classification_model_config)
         self.num_train_epochs = num_train_epochs
         self.batch_size = batch_size
@@ -132,7 +135,7 @@ def _initialize_model(self) -> Any:  # noqa: ANN401
         label2id = {i: i for i in range(self._n_classes)}
         id2label = {i: i for i in range(self._n_classes)}
 
-        return AutoModelForSequenceClassification.from_pretrained(
+        return self._AutoModelForSequenceClassification.from_pretrained(
             self.classification_model_config.model_name,
             trust_remote_code=self.classification_model_config.trust_remote_code,
             num_labels=self._n_classes,
@@ -148,7 +151,7 @@ def fit(
     ) -> None:
         self._validate_task(labels)
 
-        self._tokenizer = AutoTokenizer.from_pretrained(self.classification_model_config.model_name)  # type: ignore[no-untyped-call]
+        self._tokenizer = self._AutoTokenizer.from_pretrained(self.classification_model_config.model_name)  # type: ignore[no-untyped-call]
         self._model = self._initialize_model()
         tokenized_dataset = self._get_tokenized_dataset(utterances, labels)
         self._train(tokenized_dataset)
@@ -162,7 +165,7 @@ def _train(self, tokenized_dataset: DatasetDict) -> None:
             tokenized_dataset: output from :py:meth:`BertScorer._get_tokenized_dataset`
         """
         with tempfile.TemporaryDirectory() as tmp_dir:
-            training_args = TrainingArguments(
+            training_args = self._TrainingArguments(
                 output_dir=tmp_dir,
                 num_train_epochs=self.num_train_epochs,
                 per_device_train_batch_size=self.batch_size,
@@ -181,27 +184,27 @@ def _train(self, tokenized_dataset: DatasetDict) -> None:
                 load_best_model_at_end=self.early_stopping_config.metric is not None,
             )
 
-            trainer = Trainer(
+            trainer = self._Trainer(
                 model=self._model,
                 args=training_args,
                 train_dataset=tokenized_dataset["train"],
                 eval_dataset=tokenized_dataset["validation"],
                 processing_class=self._tokenizer,
-                data_collator=DataCollatorWithPadding(tokenizer=self._tokenizer),
+                data_collator=self._DataCollatorWithPadding(tokenizer=self._tokenizer),
                 compute_metrics=self._get_compute_metrics(),
                 callbacks=self._get_trainer_callbacks(),
             )
             if not self.print_progress:
-                trainer.remove_callback(PrinterCallback)
-                trainer.remove_callback(ProgressCallback)
+                trainer.remove_callback(self._PrinterCallback)
+                trainer.remove_callback(self._ProgressCallback)
 
             trainer.train()
 
-    def _get_trainer_callbacks(self) -> list[TrainerCallback]:
-        res: list[TrainerCallback] = []
+    def _get_trainer_callbacks(self) -> list["TrainerCallback"]:
+        res: list["TrainerCallback"] = []
         if self.early_stopping_config.metric is not None:
             res.append(
-                EarlyStoppingCallback(
+                self._EarlyStoppingCallback(
                     early_stopping_patience=self.early_stopping_config.patience,
                     early_stopping_threshold=self.early_stopping_config.threshold,
                 )
@@ -235,7 +238,7 @@ def tokenize_function(examples: dict[str, Any]) -> dict[str, Any]:
 
         return dataset.map(tokenize_function, batched=True, batch_size=self.batch_size)
 
-    def _get_compute_metrics(self) -> Callable[[EvalPrediction], dict[str, float]] | None:
+    def _get_compute_metrics(self) -> Callable[["EvalPrediction"], dict[str, float]] | None:
         """Construct callable for computing metrics during transformer training.
 
         The result of this function is supposed to pass to :py:class:`transformers.Trainer`.
@@ -246,7 +249,7 @@ def _get_compute_metrics(self) -> Callable[[EvalPrediction], dict[str, float]] |
         metric_name = self.early_stopping_config.metric
         metric_fn = (SCORING_METRICS_MULTILABEL | SCORING_METRICS_MULTICLASS)[metric_name]
 
-        def compute_metrics(output: EvalPrediction) -> dict[str, float]:
+        def compute_metrics(output: "EvalPrediction") -> dict[str, float]:
             return {
                 metric_name: metric_fn(output.label_ids.tolist(), output.predictions.tolist())  # type: ignore[union-attr]
             }

Original file line number	Diff line number	Diff line change
`@@ -52,6 +52,7 @@ dependencies = [`
`52`	`52`	`[project.optional-dependencies]`
`53`	`53`	`catboost = ["catboost (>=1.2.8,<2.0.0)"]`
`54`	`54`	`peft = ["peft (>= 0.10.0, !=0.15.0, !=0.15.1, <1.0.0)"]`
	`55`	`+transformers = ["transformers[torch] (>=4.49.0,<5.0.0)"]`
`55`	`56`	`dspy = [`
`56`	`57`	`"dspy (>=2.6.5,<3.0.0)",`
`57`	`58`	`]`