Add optional argument original_output to vLLM's generate function

outlines/models/vllm.py

@@ -50,6 +50,7 @@ def generate(
         *,
         sampling_params: Optional["SamplingParams"] = None,
         use_tqdm: bool = True,
+        original_output: bool = False,
     ):
         """Generate text using vLLM.
 
@@ -74,6 +75,9 @@ def generate(
             vLLM documentation for more details: https://docs.vllm.ai/en/latest/dev/sampling_params.html.
         use_tqdm
             A boolean in order to display progress bar while inferencing
+        original_output
+            A boolean, if True then returns the original (full) output of the
+            vLLM model
 
         Returns
         -------
@@ -82,6 +86,8 @@ def generate(
         this is a batch with several sequences but only one sample the list is
         of shape `(n_batch)`. If there is only one sequence and one sample, a
         string is returned.
+        If original_output is True, then this function returns the original
+        (full) output of the vLLM model.
 
         """
         from vllm.sampling_params import SamplingParams
@@ -134,6 +140,10 @@ def generate(
             lora_request=self.lora_request,
             use_tqdm=use_tqdm,
         )
+
+        if original_output:
+            return results
+
         results = [[sample.text for sample in batch.outputs] for batch in results]
 
         batch_size = len(results)

tests/generate/test_integration_vllm.py

@@ -4,6 +4,7 @@
 import pytest
 import torch
 from pydantic import BaseModel, constr
+from vllm.outputs import RequestOutput
 from vllm.sampling_params import SamplingParams
 
 import outlines.generate as generate
@@ -42,15 +43,31 @@ def test_vllm_generation_api(model, generator_type, params):
     res = generator("test", stop_at=[".", "ab"])
     assert isinstance(res, str)
 
+    res = generator("test", original_output=True)
+    assert isinstance(res, list)
+    assert len(res) == 1
+    assert isinstance(res[0], RequestOutput)
+
     res1 = generator("test", seed=1)
     res2 = generator("test", seed=1)
     assert isinstance(res1, str)
     assert isinstance(res2, str)
     assert res1 == res2
 
+    res1 = generator("test", seed=1, original_output=True)
+    res2 = generator("test", seed=1)
+    assert isinstance(res1[0], RequestOutput)
+    assert isinstance(res2, str)
+    text1 = [sample.text for sample in res1[0].outputs]
+    assert len(text1) == 1
+    assert text1[0] == res2
+
     res = generator(["test", "test1"])
     assert len(res) == 2
 
+    res = generator(["test", "test1"], original_output=True)
+    assert len(res) == 2
+
 
 def test_vllm_sampling_params(model):
     generator = generate.text(model)