make sentence-transformers an optional dependency

Author: voorhs

Makefile

@@ -8,7 +8,7 @@ sh = uv run --no-sync --frozen
 .PHONY: install
 install:
 	rm -rf uv.lock 
-	uv sync --all-groups --extra catboost --extra peft
+	uv sync --all-groups --extra catboost --extra peft --extra sentence-transformers --extra transformers
 
 .PHONY: test
 test:

pyproject.toml

@@ -31,7 +31,7 @@ classifiers=[
 ]
 requires-python = ">=3.10,<3.13"
 dependencies = [
-    "sentence-transformers (>=3,<4)",
+    "torch (>=2.0.0,<3.0.0)",
     "scikit-learn (>=1.5,<2.0)",
     "iterative-stratification (>=0.1.9)",
     "appdirs (>=1.4,<2.0)",
@@ -43,7 +43,6 @@ dependencies = [
     "datasets (>=3.2.0,<4.0.0)",
     "xxhash (>=3.5.0,<4.0.0)",
     "python-dotenv (>=1.0.1,<2.0.0)",
-    "transformers[torch] (>=4.49.0,<5.0.0)",
     "aiometer (>=1.0.0,<2.0.0)",
     "aiofiles (>=24.1.0,<25.0.0)",
     "threadpoolctl (>=3.0.0,<4.0.0)",
@@ -52,7 +51,8 @@ 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)"]
+transformers = ["transformers (>=4.49.0,<5.0.0)"]
+sentence-transformers = ["sentence-transformers (>=3,<4)"]
 dspy = [
     "dspy (>=2.6.5,<3.0.0)",
 ]

src/autointent/_wrappers/embedder/sentence_transformers.py

@@ -10,9 +10,6 @@
 import numpy.typing as npt
 import torch
 from datasets import Dataset
-from sentence_transformers import SentenceTransformer, SentenceTransformerTrainer, SentenceTransformerTrainingArguments
-from sentence_transformers.losses import BatchAllTripletLoss
-from sentence_transformers.training_args import BatchSamplers
 from sklearn.model_selection import train_test_split
 
 from autointent._hash import Hasher
@@ -25,6 +22,7 @@
 from .utils import get_embeddings_path
 
 if TYPE_CHECKING:
+    from sentence_transformers import SentenceTransformer
     from transformers import TrainerCallback
 
 logger = logging.getLogger(__name__)
@@ -51,6 +49,7 @@ class SentenceTransformerEmbeddingBackend(BaseEmbeddingBackend):
     """SentenceTransformer-based embedding backend implementation."""
 
     supports_training: bool = True
+    _model: "SentenceTransformer | None"
 
     def __init__(self, config: SentenceTransformerEmbeddingConfig) -> None:
         """Initialize the SentenceTransformer backend.
@@ -59,7 +58,7 @@ def __init__(self, config: SentenceTransformerEmbeddingConfig) -> None:
             config: Configuration for SentenceTransformer embeddings.
         """
         self.config = config
-        self._model: SentenceTransformer | None = None
+        self._model = None
         self._trained: bool = False
 
     def clear_ram(self) -> None:
@@ -71,10 +70,12 @@ def clear_ram(self) -> None:
             self._model = None
             torch.cuda.empty_cache()
 
-    def _load_model(self) -> SentenceTransformer:
+    def _load_model(self) -> "SentenceTransformer":
         """Load sentence transformers model to device."""
         if self._model is None:
-            res = SentenceTransformer(
+            # Lazy import sentence-transformers
+            st = require("sentence_transformers", extra="sentence-transformers")
+            res = st.SentenceTransformer(
                 self.config.model_name,
                 device=self.config.device,
                 prompts=self.config.get_prompt_config(),
@@ -231,16 +232,17 @@ def train(self, utterances: list[str], labels: ListOfLabels, config: EmbedderFin
 
         model = self._load_model()
 
+        # Lazy import sentence-transformers training components (only needed for fine-tuning)
+        st = require("sentence_transformers", extra="sentence-transformers")
+        transformers = require("transformers", extra="transformers")
+
         x_train, x_val, y_train, y_val = train_test_split(utterances, labels, test_size=config.val_fraction)
         tr_ds = Dataset.from_dict({"text": x_train, "label": y_train})
         val_ds = Dataset.from_dict({"text": x_val, "label": y_val})
 
-        loss = BatchAllTripletLoss(model=model, margin=config.margin)
+        loss = st.losses.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(
+            args = st.SentenceTransformerTrainingArguments(
                 save_strategy="epoch",
                 save_total_limit=1,
                 output_dir=tmp_dir,
@@ -251,7 +253,7 @@ def train(self, utterances: list[str], labels: ListOfLabels, config: EmbedderFin
                 warmup_ratio=config.warmup_ratio,
                 fp16=config.fp16,
                 bf16=config.bf16,
-                batch_sampler=BatchSamplers.NO_DUPLICATES,
+                batch_sampler=st.training_args.BatchSamplers.NO_DUPLICATES,
                 metric_for_best_model="eval_loss",
                 load_best_model_at_end=True,
                 eval_strategy="epoch",
@@ -263,7 +265,7 @@ def train(self, utterances: list[str], labels: ListOfLabels, config: EmbedderFin
                     early_stopping_threshold=config.early_stopping_threshold,
                 )
             ]
-            trainer = SentenceTransformerTrainer(
+            trainer = st.SentenceTransformerTrainer(
                 model=model,
                 args=args,
                 train_dataset=tr_ds,

src/autointent/_wrappers/ranker.py

@@ -10,19 +10,22 @@
 import logging
 from pathlib import Path
 from random import shuffle
-from typing import Any, Literal, TypedDict
+from typing import TYPE_CHECKING, Any, Literal, TypedDict
 
 import joblib
 import numpy as np
 import numpy.typing as npt
-import sentence_transformers as st
 import torch
 from sklearn.linear_model import LogisticRegressionCV
 from torch import nn
 
+from autointent._utils import require
 from autointent.configs import CrossEncoderConfig
 from autointent.custom_types import ListOfLabels, RerankedItem
 
+if TYPE_CHECKING:
+    import sentence_transformers as st
+
 logger = logging.getLogger(__name__)
 
 
@@ -95,7 +98,7 @@ class Ranker:
     _metadata_file_name = "metadata.json"
     _classifier_file_name = "classifier.joblib"
     config: CrossEncoderConfig
-    cross_encoder: st.CrossEncoder
+    cross_encoder: "st.CrossEncoder"
 
     def __init__(
         self,
@@ -110,12 +113,15 @@ def __init__(
             classifier_head: Optional pre-trained classifier head
             output_range: Range of the output probabilities ([0, 1] for sigmoid, [-1, 1] for tanh)
         """
+        # Lazy import sentence-transformers
+        st = require("sentence_transformers", extra="sentence-transformers")
+
         self.config = CrossEncoderConfig.from_search_config(cross_encoder_config)
         self.cross_encoder = st.CrossEncoder(
             self.config.model_name,
             trust_remote_code=self.config.trust_remote_code,
             device=self.config.device,
-            max_length=self.config.tokenizer_config.max_length,  # type: ignore[arg-type]
+            max_length=self.config.tokenizer_config.max_length,
         )
         self._train_head = False
         self._clf = classifier_head