[Refactor] vLLMWrapper class

ghstack-source-id: 21ed901c0428451312cad5fb85e149a9dd694819
Pull Request resolved: #2870

Author: vmoens

test/test_actors.py

@@ -23,11 +23,11 @@
 from torchrl.envs import LLMEnv
 from torchrl.modules import (
     from_hf_transformers,
-    from_vllm,
     MLP,
     SafeModule,
     TanhDelta,
     TanhNormal,
+    vLLMWrapper,
 )
 from torchrl.modules.tensordict_module.actors import (
     _process_action_space_spec,
@@ -1031,7 +1031,7 @@ def test_from_vllm(
         torch.manual_seed(0)
 
         model = vllm_instance
-        m = from_vllm(
+        m = vLLMWrapper(
             model,
             from_text=from_text,
             generate=generate,
@@ -1207,14 +1207,14 @@ def test_from_vllm_logprobs(
         torch.manual_seed(0)
 
         model = vllm_instance
-        m_generate = from_vllm(
+        m_generate = vLLMWrapper(
             model,
             from_text=from_text,
             generate=True,
             return_log_probs=True,
             pad_output=pad_output,
         )
-        m_logprobs = from_vllm(
+        m_logprobs = vLLMWrapper(
             model, from_text=from_text, generate=False, pad_output=pad_output
         )
         self._check_lps(
@@ -1264,7 +1264,7 @@ def _check_lps(
     @pytest.mark.parametrize("use_tensorclass", [True, False])
     def test_vllm_batch_run(self, pad, generate, use_tensorclass, vllm_instance):
         # Test generate - padding combinations
-        policy = from_vllm(
+        policy = vLLMWrapper(
             vllm_instance,
             from_text=True,
             generate=generate,
@@ -1331,7 +1331,7 @@ def test_vllm_batch_run(self, pad, generate, use_tensorclass, vllm_instance):
             assert isinstance(tokens, list)
 
     def test_vllm_collection(self, vllm_instance):
-        policy = from_vllm(
+        policy = vLLMWrapper(
             vllm_instance,
             from_text=True,
             generate=True,

torchrl/modules/__init__.py

@@ -93,7 +93,7 @@
 )
 from .utils import get_primers_from_module
 from .planners import CEMPlanner, MPCPlannerBase, MPPIPlanner  # usort:skip
-from .llm import from_hf_transformers, from_vllm
+from .llm import from_hf_transformers, vLLMWrapper
 
 __all__ = [
     "Actor",
@@ -178,7 +178,7 @@
     "WorldModelWrapper",
     "distributions_maps",
     "from_hf_transformers",
-    "from_vllm",
+    "vLLMWrapper",
     "get_primers_from_module",
     "recurrent_mode",
     "reset_noise",

torchrl/modules/llm/__init__.py

@@ -5,6 +5,7 @@
 
 from .common import CategoricalSequential
 from .transformers_policy import from_hf_transformers
-from .vllm_policy import from_vllm
 
-__all__ = ["from_hf_transformers", "from_vllm", "CategoricalSequential"]
+from .vllm_wrapper import vLLMWrapper
+
+__all__ = ["from_hf_transformers", "vLLMWrapper", "CategoricalSequential"]

torchrl/modules/llm/common.py

@@ -7,14 +7,14 @@
 from tensordict import NestedKey, TensorDictBase
 from tensordict.nn import (
     ProbabilisticTensorDictModule,
-    ProbabilisticTensorDictSequential,
+    TensorDictModuleBase,
     TensorDictSequential,
 )
 from torch import distributions as D
 from torch.distributions import Categorical
 
 
-class CategoricalSequential(ProbabilisticTensorDictSequential):
+class CategoricalSequential(TensorDictModuleBase):
     """A ProbabilisticTensorDictSequential subclass meant to work with LLMs.
 
     .. seealso:: :class:`~tensordict.nn.ProbabilisticTensorDictSequential` class.

torchrl/modules/llm/transformers_policy.py

@@ -181,7 +181,7 @@ def from_hf_transformers(
         >>> print(output_data.text_response)
         [' heritageillon rooft rooft Pear Tes grantingalde 58ocrocrocrocrcubecubecubecubecubecubecube']
 
-    .. seealso:: :func:`~torchrl.modules.from_vllm` for a similar interface using the vLLM library.
+    .. seealso:: :func:`~torchrl.modules.vLLMWrapper` for a similar interface using the vLLM library.
 
     """
     # TODO: Seq should have a return_log_prob and be of ProbabilisticTDSequential type for instance checks