Move VLLMClient to generation module (huggingface#4928)

Author: albertvillanova

tests/test_vllm_client_server.py

@@ -19,7 +19,7 @@
 from transformers import AutoModelForCausalLM
 from transformers.testing_utils import torch_device
 
-from trl.extras.vllm_client import VLLMClient
+from trl.generation.vllm_client import VLLMClient
 from trl.import_utils import is_vllm_available
 from trl.scripts.vllm_serve import chunk_list
 

trl/experimental/gold/gold_trainer.py

@@ -48,7 +48,7 @@
 
 from ...data_utils import is_conversational, maybe_convert_to_chatml, pack_dataset, truncate_dataset
 from ...extras.profiling import profiling_decorator
-from ...extras.vllm_client import VLLMClient
+from ...generation.vllm_client import VLLMClient
 from ...import_utils import is_vllm_available
 from ...models import prepare_deepspeed
 from ...models.utils import unwrap_model_for_generation

trl/experimental/online_dpo/online_dpo_trainer.py

@@ -53,7 +53,7 @@
 
 from ...data_utils import apply_chat_template, is_conversational, maybe_apply_chat_template
 from ...extras.profiling import profiling_context
-from ...extras.vllm_client import VLLMClient
+from ...generation.vllm_client import VLLMClient
 from ...import_utils import is_vllm_available
 from ...models.utils import create_reference_model, prepare_deepspeed, prepare_fsdp, unwrap_model_for_generation
 from ...trainer.base_trainer import BaseTrainer

trl/extras/vllm_client.py trl/generation/vllm_client.pytrl/extras/vllm_client.py renamed to trl/generation/vllm_client.py

@@ -28,7 +28,7 @@
 from transformers import is_torch_xpu_available
 from urllib3.util.retry import Retry
 
-from ..import_utils import is_requests_available, is_vllm_ascend_available, is_vllm_available
+from import_utils import is_requests_available, is_vllm_ascend_available, is_vllm_available
 
 
 if is_requests_available():
@@ -88,7 +88,7 @@ class VLLMClient:
         Use the client to generate completions and update model weights:
 
         ```python
-        >>> from trl.extras.vllm_client import VLLMClient
+        >>> from trl.generation.vllm_client import VLLMClient
 
         >>> client = VLLMClient()
         >>> client.generate(["Hello, AI!", "Tell me a joke"])

trl/generation/vllm_generation.py

@@ -29,9 +29,9 @@
 
 from ..data_utils import apply_chat_template, is_conversational, prepare_multimodal_messages_vllm
 from ..extras.profiling import ProfilingContext, profiling_decorator
-from ..extras.vllm_client import VLLMClient
 from ..import_utils import is_vllm_available
 from ..trainer.utils import ensure_master_addr_port
+from .vllm_client import VLLMClient
 
 
 if TYPE_CHECKING: