Make vLLM and transformers optional dependencies

README.md

@@ -61,6 +61,22 @@ conda install -c conda-forge openjdk=21 maven -y
 pip install "rank-llm[pyserini]"
 ```
 
+## Core Optional Dependencies
+
+### Install vLLM Support (Optional)
+For local model inference using vLLM:
+```bash
+pip install -e .[vllm]      # local installation for development
+pip install rank-llm[vllm]  # or pip installation
+```
+
+### Install Transformers Support (Optional)
+For T5-based models (DuoT5, MonoT5, etc.):
+```bash
+pip install -e .[transformers]      # local installation for development
+pip install rank-llm[transformers]  # or pip installation
+```
+
 ## Install [all] Dependencies
 ```bash
 pip install -e .[all]      # local installation for development

docs/optional_dependencies.md

@@ -0,0 +1,82 @@
+# Optional Dependencies in rank_llm
+
+rank_llm now supports optional dependencies to allow for lighter installations based on your specific needs.
+
+## Installation Options
+
+### Core Installation
+Install just the core functionality without heavy ML libraries:
+```bash
+pip install rank_llm
+```
+
+This includes the basic ranking functionality and API integrations, but excludes:
+- vLLM for local model inference
+- transformers library for T5-based models
+
+### With vLLM Support
+For local model inference using vLLM:
+```bash
+pip install rank_llm[vllm]
+```
+
+This enables:
+- `VllmHandler` for local model inference
+- `RankListwiseOSLLM` and other vLLM-dependent models
+
+### With Transformers Support
+For T5-based models and transformers functionality:
+```bash
+pip install rank_llm[transformers]
+```
+
+This enables:
+- `DuoT5` pairwise ranking model
+- `MonoT5` pointwise ranking model
+- `RankFiD` listwise models
+- All T5-based inference handlers
+
+### Full Installation
+Install all optional dependencies:
+```bash
+pip install rank_llm[all]
+```
+
+This includes vLLM, transformers, and all other optional features.
+
+## Error Handling
+
+When you try to use functionality that requires missing optional dependencies, you'll get helpful error messages:
+
+```python
+from rank_llm.rerank.vllm_handler import VllmHandler
+
+# If vLLM is not installed:
+handler = VllmHandler(...)
+# ImportError: vLLM is not installed. Please install it with: pip install rank_llm[vllm]
+```
+
+```python
+from rank_llm.rerank.pairwise.duot5 import DuoT5
+
+# If transformers is not installed:
+model = DuoT5(...)
+# ImportError: transformers is not installed. Please install it with: pip install rank_llm[transformers]
+```
+
+## Migration Guide
+
+If you were previously using rank_llm and now get import errors, you likely need to install the optional dependencies:
+
+1. **For vLLM users**: Run `pip install rank_llm[vllm]`
+2. **For T5 model users**: Run `pip install rank_llm[transformers]`
+3. **For both**: Run `pip install rank_llm[all]`
+
+## Development
+
+When developing rank_llm, install all dependencies:
+```bash
+pip install -e .[all]
+```
+
+This ensures you can test all functionality locally.

pyproject.toml

@@ -27,6 +27,12 @@ requires-python = ">= 3.11"
 dependencies = {file = ["requirements.txt"]}
 
 [project.optional-dependencies]
+vllm = [
+    "vllm>=0.4.0"
+]
+transformers = [
+    "transformers>=4.40.1"
+]
 sglang = [
     "sglang[all]~=0.4.0"
 ]
@@ -50,7 +56,9 @@ training = [
 ]
 all = [
   "rank-llm[genai]",
-  "rank-llm[pyserini]"
+  "rank-llm[pyserini]",
+  "rank-llm[vllm]",
+  "rank-llm[transformers]"
 ]
 
 [project.urls]

requirements.txt

@@ -1,10 +1,8 @@
 tqdm>=4.66.2
 openai>=1.23.6
 tiktoken>=0.6.0
-transformers>=4.40.1
 python-dotenv>=1.0.1
 faiss-cpu>=1.8.0
 ftfy>=6.2.0
 dacite>=1.8.1
-vllm>=0.4.0
 pandas>=1.4.0

src/rank_llm/rerank/listwise/lit5/model.py

@@ -3,7 +3,14 @@
 
 import torch
 from torch import nn
-from transformers.models.t5.modeling_t5 import T5ForConditionalGeneration, T5Stack
+
+try:
+    from transformers.models.t5.modeling_t5 import T5ForConditionalGeneration, T5Stack
+    TRANSFORMERS_AVAILABLE = True
+except ImportError:
+    raise ImportError(
+        "transformers is required for the lit5 model. Please install it with: pip install rank_llm[transformers]"
+    )
 
 from .modeling_t5 import (
     T5ForConditionalGeneration as T5ConditionalGenerationCrossAttentionScore,

src/rank_llm/rerank/listwise/lit5/modeling_t5.py

@@ -24,30 +24,37 @@
 from torch import nn
 from torch.nn import CrossEntropyLoss
 from torch.utils.checkpoint import checkpoint
-from transformers.activations import ACT2FN
-from transformers.file_utils import (
-    DUMMY_INPUTS,
-    DUMMY_MASK,
-    add_start_docstrings,
-    add_start_docstrings_to_model_forward,
-    is_torch_fx_proxy,
-    replace_return_docstrings,
-)
-from transformers.generation import GenerationMixin
-from transformers.modeling_outputs import (
-    BaseModelOutput,
-    BaseModelOutputWithPastAndCrossAttentions,
-    Seq2SeqLMOutput,
-    Seq2SeqModelOutput,
-)
-from transformers.modeling_utils import (
-    PreTrainedModel,
-    find_pruneable_heads_and_indices,
-    prune_linear_layer,
-)
-from transformers.models.t5.configuration_t5 import T5Config
-from transformers.utils import logging
-from transformers.utils.model_parallel_utils import assert_device_map, get_device_map
+
+try:
+    from transformers.activations import ACT2FN
+    from transformers.file_utils import (
+        DUMMY_INPUTS,
+        DUMMY_MASK,
+        add_start_docstrings,
+        add_start_docstrings_to_model_forward,
+        is_torch_fx_proxy,
+        replace_return_docstrings,
+    )
+    from transformers.generation import GenerationMixin
+    from transformers.modeling_outputs import (
+        BaseModelOutput,
+        BaseModelOutputWithPastAndCrossAttentions,
+        Seq2SeqLMOutput,
+        Seq2SeqModelOutput,
+    )
+    from transformers.modeling_utils import (
+        PreTrainedModel,
+        find_pruneable_heads_and_indices,
+        prune_linear_layer,
+    )
+    from transformers.models.t5.configuration_t5 import T5Config
+    from transformers.utils import logging
+    from transformers.utils.model_parallel_utils import assert_device_map, get_device_map
+    TRANSFORMERS_AVAILABLE = True
+except ImportError:
+    raise ImportError(
+        "transformers is required for the lit5 model. Please install it with: pip install rank_llm[transformers]"
+    )
 
 logger = logging.get_logger(__name__)
 

src/rank_llm/rerank/listwise/rank_fid.py

@@ -3,7 +3,13 @@
 
 import torch
 from tqdm import tqdm
-from transformers import T5Tokenizer
+
+try:
+    from transformers import T5Tokenizer
+    TRANSFORMERS_AVAILABLE = True
+except ImportError:
+    TRANSFORMERS_AVAILABLE = False
+    T5Tokenizer = None
 
 from rank_llm.data import Request, Result
 from rank_llm.rerank.listwise.listwise_rankllm import ListwiseRankLLM
@@ -45,6 +51,11 @@ def __init__(
         """
         Creates instance of the RankFiDDistill class, a specialized version of RankLLM designed from Lit5-Distill.
         """
+        if not TRANSFORMERS_AVAILABLE:
+            raise ImportError(
+                "transformers is not installed. Please install it with: pip install rank_llm[transformers]"
+            )
+
         super().__init__(
             model=model,
             context_size=context_size,
@@ -269,6 +280,11 @@ def __init__(
         precision: str = "bfloat16",
         device: str = "cuda",
     ) -> None:
+        if not TRANSFORMERS_AVAILABLE:
+            raise ImportError(
+                "transformers is not installed. Please install it with: pip install rank_llm[transformers]"
+            )
+
         super().__init__(
             model=model,
             context_size=context_size,

src/rank_llm/rerank/listwise/rank_listwise_os_llm.py

@@ -4,13 +4,19 @@
 import unicodedata
 from concurrent.futures import ThreadPoolExecutor
 from importlib.resources import files
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 import torch
-import vllm
 from ftfy import fix_text
 from tqdm import tqdm
 
+try:
+    import vllm
+    VLLM_AVAILABLE = True
+except ImportError:
+    VLLM_AVAILABLE = False
+    vllm = None
+
 from rank_llm.data import Request, Result
 from rank_llm.rerank.rankllm import PromptMode
 from rank_llm.rerank.vllm_handler import VllmHandler
@@ -225,7 +231,7 @@ def _evaluate_logits(
 
     def _get_logits_single_digit(
         self,
-        output: vllm.RequestOutput,
+        output: Union["vllm.RequestOutput", Any],
         effective_location: int = 1,
         total: Tuple[int, int] = (1, 9),
     ):

src/rank_llm/rerank/pairwise/duot5.py

@@ -4,8 +4,15 @@
 from importlib.resources import files
 from typing import List, Optional, Tuple
 
-from transformers import T5ForConditionalGeneration, T5Tokenizer
-from transformers.generation import GenerationConfig
+try:
+    from transformers import T5ForConditionalGeneration, T5Tokenizer
+    from transformers.generation import GenerationConfig
+    TRANSFORMERS_AVAILABLE = True
+except ImportError:
+    TRANSFORMERS_AVAILABLE = False
+    T5ForConditionalGeneration = None
+    T5Tokenizer = None
+    GenerationConfig = None
 
 from rank_llm.data import Result
 from rank_llm.rerank.pairwise.pairwise_rankllm import PairwiseRankLLM
@@ -28,6 +35,11 @@ def __init__(
         device: str = "cuda",
         batch_size: int = 32,
     ):
+        if not TRANSFORMERS_AVAILABLE:
+            raise ImportError(
+                "transformers is not installed. Please install it with: pip install rank_llm[transformers]"
+            )
+
         super().__init__(
             model=model,
             context_size=context_size,

src/rank_llm/rerank/pairwise/pairwise_inference_handler.py

@@ -1,7 +1,12 @@
 import re
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Union
 
-from transformers import T5Tokenizer
+try:
+    from transformers import T5Tokenizer
+    TRANSFORMERS_AVAILABLE = True
+except ImportError:
+    TRANSFORMERS_AVAILABLE = False
+    T5Tokenizer = None
 
 from rank_llm.data import Result, TemplateSectionConfig
 from rank_llm.rerank.inference_handler import BaseInferenceHandler
@@ -96,7 +101,7 @@ def _generate_body(
         index1: int,
         index2: int,
         single_doc_max_token: int,
-        tokenizer: T5Tokenizer,
+        tokenizer: Union["T5Tokenizer", Any],
     ) -> str:
         doc1_raw = self._convert_doc_to_prompt_content(
             result.candidates[index1].doc, max_length=single_doc_max_token

src/rank_llm/rerank/pointwise/monot5.py

@@ -3,8 +3,15 @@
 from importlib.resources import files
 from typing import List, Optional, Tuple
 
-from transformers import T5ForConditionalGeneration, T5Tokenizer
-from transformers.generation import GenerationConfig
+try:
+    from transformers import T5ForConditionalGeneration, T5Tokenizer
+    from transformers.generation import GenerationConfig
+    TRANSFORMERS_AVAILABLE = True
+except ImportError:
+    TRANSFORMERS_AVAILABLE = False
+    T5ForConditionalGeneration = None
+    T5Tokenizer = None
+    GenerationConfig = None
 
 from rank_llm.data import Result
 from rank_llm.rerank.pointwise.pointwise_rankllm import PointwiseRankLLM
@@ -27,6 +34,11 @@ def __init__(
         device: str = "cuda",
         batch_size: int = 32,
     ):
+        if not TRANSFORMERS_AVAILABLE:
+            raise ImportError(
+                "transformers is not installed. Please install it with: pip install rank_llm[transformers]"
+            )
+
         super().__init__(
             model=model,
             context_size=context_size,

src/rank_llm/rerank/pointwise/pointwise_inference_handler.py

@@ -1,7 +1,12 @@
 import re
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Union
 
-from transformers import T5Tokenizer
+try:
+    from transformers import T5Tokenizer
+    TRANSFORMERS_AVAILABLE = True
+except ImportError:
+    TRANSFORMERS_AVAILABLE = False
+    T5Tokenizer = None
 
 from rank_llm.data import Result, TemplateSectionConfig
 from rank_llm.rerank.inference_handler import BaseInferenceHandler
@@ -86,7 +91,7 @@ def _generate_body(
         result: Result,
         index: int,
         max_doc_tokens: int,
-        tokenizer: T5Tokenizer,
+        tokenizer: Union["T5Tokenizer", Any],
     ) -> str:
         query = self._replace_number(result.query.text)
         doc_raw = self._convert_doc_to_prompt_content(