outputs.py 19.4 KB
Newer Older
1
import time
2
from dataclasses import dataclass
3
from typing import Dict, Generic, List, MutableSequence, Optional
4
5
from typing import Sequence as GenericSequence
from typing import Union
6

7
import torch
8
from typing_extensions import TypeVar, deprecated
9

10
from vllm.lora.request import LoRARequest
11
from vllm.multimodal.inputs import MultiModalPlaceholderDict
12
from vllm.sampling_params import RequestOutputKind
13
from vllm.sequence import (PromptLogprobs, RequestMetrics, SampleLogprobs,
14
                           SequenceGroup, SequenceGroupBase, SequenceStatus)
15
16


17
@dataclass
18
class CompletionOutput:
Zhuohan Li's avatar
Zhuohan Li committed
19
20
21
22
23
24
25
26
27
28
29
    """The output data of one completion output of a request.

    Args:
        index: The index of the output in the request.
        text: The generated output text.
        token_ids: The token IDs of the generated output text.
        cumulative_logprob: The cumulative log probability of the generated
            output text.
        logprobs: The log probabilities of the top probability words at each
            position if the logprobs are requested.
        finish_reason: The reason why the sequence is finished.
30
31
32
        stop_reason: The stop string or token id that caused the completion
            to stop, None if the completion finished for some other reason
            including encountering the EOS token.
33
        lora_request: The LoRA request that was used to generate the output.
Zhuohan Li's avatar
Zhuohan Li committed
34
    """
35

36
37
    index: int
    text: str
38
    token_ids: GenericSequence[int]
39
    cumulative_logprob: Optional[float]
40
41
42
43
    logprobs: Optional[SampleLogprobs]
    finish_reason: Optional[str] = None
    stop_reason: Union[int, str, None] = None
    lora_request: Optional[LoRARequest] = None
Zhuohan Li's avatar
Zhuohan Li committed
44
45
46

    def finished(self) -> bool:
        return self.finish_reason is not None
47
48

    def __repr__(self) -> str:
49
50
        return (f"CompletionOutput(index={self.index}, "
                f"text={self.text!r}, "
51
                f"token_ids={self.token_ids}, "
52
                f"cumulative_logprob={self.cumulative_logprob}, "
53
                f"logprobs={self.logprobs}, "
54
55
                f"finish_reason={self.finish_reason}, "
                f"stop_reason={self.stop_reason})")
56
57


58
@dataclass
59
60
class PoolingOutput:
    """The output data of one pooling output of a request.
61
62

    Args:
63
        data: The extracted hidden states.
64
    """
65
    data: torch.Tensor
66
67

    def __repr__(self) -> str:
68
69
70
71
72
73
74
        return (f"PoolingOutput(data={self.data})")

    def __eq__(self, other: object) -> bool:
        return (isinstance(other, self.__class__) and bool(
            (self.data == other.data).all()))

    @property
75
76
77
78
    @deprecated("`LLM.encode()` now stores raw outputs in the `data` "
                "attribute. To return embeddings, use `LLM.embed()`. "
                "To return class probabilities, use `LLM.classify()` "
                "and access the `probs` attribute. ")
79
80
    def embedding(self) -> list[float]:
        return self.data.tolist()
81
82


83
class RequestOutput:
84
    """The output data of a completion request to the LLM.
Zhuohan Li's avatar
Zhuohan Li committed
85
86
87
88

    Args:
        request_id: The unique ID of the request.
        prompt: The prompt string of the request.
89
90
                For encoder/decoder models, this is the
                decoder input prompt.
Zhuohan Li's avatar
Zhuohan Li committed
91
        prompt_token_ids: The token IDs of the prompt.
92
93
                          For encoder/decoder models, this is the
                          decoder input prompt token ids.
lots-o's avatar
lots-o committed
94
        prompt_logprobs: The log probabilities to return per prompt token.
Zhuohan Li's avatar
Zhuohan Li committed
95
        outputs: The output sequences of the request.
96
        finished: Whether the whole request is finished.
97
        metrics: Metrics associated with the request.
98
        lora_request: The LoRA request that was used to generate the output.
99
100
101
102
103
        encoder_prompt: The encoder prompt string of the request.
                        None if decoder-only.
        encoder_prompt_token_ids: The token IDs of the encoder prompt.
                                  None if decoder-only.
        num_cached_tokens: The number of tokens with prefix cache hit.
Zhuohan Li's avatar
Zhuohan Li committed
104
    """
105

106
107
    def __init__(
        self,
108
        request_id: str,
109
        prompt: Optional[str],
110
        prompt_token_ids: Optional[List[int]],
111
        prompt_logprobs: Optional[PromptLogprobs],
112
        outputs: List[CompletionOutput],
113
        finished: bool,
114
        metrics: Optional[RequestMetrics] = None,
115
        lora_request: Optional[LoRARequest] = None,
116
117
        encoder_prompt: Optional[str] = None,
        encoder_prompt_token_ids: Optional[List[int]] = None,
118
        num_cached_tokens: Optional[int] = None,
119
120
        *,
        multi_modal_placeholders: Optional[MultiModalPlaceholderDict] = None,
121
122
123
124
    ) -> None:
        self.request_id = request_id
        self.prompt = prompt
        self.prompt_token_ids = prompt_token_ids
125
        self.multi_modal_placeholders = multi_modal_placeholders or {}
126
        self.prompt_logprobs = prompt_logprobs
127
        self.outputs = outputs
128
        self.finished = finished
129
        self.metrics = metrics
130
        self.lora_request = lora_request
131
132
        self.encoder_prompt = encoder_prompt
        self.encoder_prompt_token_ids = encoder_prompt_token_ids
133
        self.num_cached_tokens = num_cached_tokens
134

135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
    @classmethod
    def new(
        cls,
        request_id: str,
        prompt: Optional[str],
        prompt_token_ids: Optional[List[int]],
        text: str,
        token_ids: List[int],
        finished: bool = False,
    ) -> "RequestOutput":
        """Initialize a new RequestOutput object."""

        # TODO: Support `n` > 1.
        completion_output = CompletionOutput(
            index=0,
            text=text,
            token_ids=token_ids,
            cumulative_logprob=None,
            logprobs=None,  # TODO
        )

        return RequestOutput(
            request_id=request_id,
            prompt=prompt,
            prompt_token_ids=prompt_token_ids,
            prompt_logprobs=None,  # TODO
            outputs=[completion_output],
            finished=finished,
        )

165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
    def add(self, next_output: "RequestOutput") -> None:
        """Merge subsequent RequestOutput into this one"""

        self.prompt = next_output.prompt
        self.prompt_token_ids = next_output.prompt_token_ids
        self.prompt_logprobs = next_output.prompt_logprobs
        self.finished |= next_output.finished

        #TODO assuming n == 1 for now
        completion = self.outputs[0]
        next_completion = next_output.outputs[0]
        completion.text += next_completion.text
        if not isinstance(completion.token_ids, MutableSequence):
            completion.token_ids = list(completion.token_ids)
        completion.token_ids.extend(next_completion.token_ids)
        if next_completion.logprobs:
            assert completion.logprobs is not None
            completion.logprobs.extend(next_completion.logprobs)
        completion.cumulative_logprob = next_completion.cumulative_logprob

185
    @classmethod
186
187
188
189
190
191
192
193
194
    def from_seq_group(
        cls, seq_group: SequenceGroup, use_cache: bool,
        seq_id_to_seq_group: Dict[str, SequenceGroupBase]
    ) -> Optional["RequestOutput"]:
        finished = seq_group.is_finished()

        if seq_group.request_id in seq_id_to_seq_group:
            group: SequenceGroupBase = seq_id_to_seq_group[
                seq_group.request_id]
195
            assembled_seq_group = group.maybe_assemble_group(seq_group)
196
197
198
199
200
201
202
            if finished:
                group.finish_seq(seq_group)
            if assembled_seq_group is None:
                return None
            return cls.from_seq_group(assembled_seq_group, use_cache,
                                      seq_id_to_seq_group)

203
204
        sampling_params = seq_group.sampling_params
        if sampling_params is None:
205
206
            raise ValueError(
                "Sampling parameters are missing for a CompletionRequest.")
207

208
209
210
211
        if sampling_params.output_kind == RequestOutputKind.FINAL_ONLY and (
                not finished):
            return None

212
213
214
215
216
217
218
219
220
221
        # Init cache (if needed)
        if use_cache and seq_group.cached_request_output is None:
            seq_group.cached_request_output = RequestOutput(  # type: ignore
                request_id="",
                prompt=None,
                prompt_token_ids=[],
                prompt_logprobs=None,
                outputs=[],
                finished=False)

222
        top_n_seqs = seq_group.get_seqs()
223

224
        # Create the outputs.
225
226
227
        # NOTE: We need omit logprobs here explicitly because the sequence
        # always has the logprobs of the sampled tokens even if the
        # logprobs are not requested.
228
229
230
231
232
233
        include_logprobs = sampling_params.logprobs is not None
        text_buffer_length = sampling_params.output_text_buffer_length
        delta = sampling_params.output_kind == RequestOutputKind.DELTA

        outputs = []
        include_prompt = True
234
235
        # num_cached_tokens should be the same for all the sequences
        num_cached_tokens = None
236
        for i, seq in enumerate(top_n_seqs):
237
238
            output_text = seq.get_output_text_to_return(
                text_buffer_length, delta)
239

240
            output_token_ids = seq.get_output_token_ids_to_return(delta)
241
242
            num_output_tokens = 1 if isinstance(output_token_ids,
                                                int) else len(output_token_ids)
243
            num_cached_tokens = seq.data.get_num_cached_tokens()
244

245
246
247
248
249
            output_logprobs = seq.output_logprobs if include_logprobs else None

            if delta:
                # Slice logprobs delta if applicable
                if output_logprobs:
250
                    output_logprobs = output_logprobs[-num_output_tokens:]
251
252
                # Don't include prompt if this is after the first output
                # containing decode token ids
253
                if include_prompt and seq.get_output_len() > num_output_tokens:
254
255
                    include_prompt = False

256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
            if use_cache:
                # Get cached output object
                cached_outputs = seq_group.cached_request_output.outputs  # type: ignore
                if i >= len(cached_outputs):
                    cached_outputs.append(
                        CompletionOutput(index=i,
                                         text="",
                                         token_ids=[],
                                         cumulative_logprob=None,
                                         logprobs=None,
                                         finish_reason=None,
                                         stop_reason=None))
                output = cached_outputs[i]

                # Init cached output object
                assert output.index == i
                output.text = output_text

                if isinstance(output_token_ids, int):
                    output.token_ids.clear()
                    output.token_ids.append(output_token_ids)
                else:
                    output.token_ids = output_token_ids

                output.cumulative_logprob = seq.get_cumulative_logprob() \
                    if include_logprobs else None
                output.logprobs = output_logprobs
                output.finish_reason = SequenceStatus.get_finished_reason(
                    seq.status)
                output.stop_reason = seq.stop_reason

            else:
                output = CompletionOutput(
289
                    top_n_seqs.index(seq), output_text, [output_token_ids]
290
                    if isinstance(output_token_ids, int) else output_token_ids,
291
292
293
                    seq.get_cumulative_logprob() if include_logprobs else None,
                    output_logprobs,
                    SequenceStatus.get_finished_reason(seq.status),
294
295
296
                    seq.stop_reason)

            outputs.append(output)
297
298

        # Every sequence in the sequence group should have the same prompt.
299
300
301
302
303
304
305
306
307
308
309
310
        if include_prompt:
            prompt = seq_group.prompt
            prompt_token_ids = seq_group.prompt_token_ids
            encoder_prompt = seq_group.encoder_prompt
            encoder_prompt_token_ids = seq_group.encoder_prompt_token_ids
            prompt_logprobs = seq_group.prompt_logprobs
        else:
            prompt = None
            prompt_token_ids = None
            encoder_prompt = None
            encoder_prompt_token_ids = None
            prompt_logprobs = None
311
312
        finished_time = time.time() if finished else None
        seq_group.set_finished_time(finished_time)
313

314
315
316
317
318
319
320
321
322
323
324
325
326
327
        init_kwargs = {
            "request_id": seq_group.request_id,
            "prompt": prompt,
            "prompt_token_ids": prompt_token_ids,
            "prompt_logprobs": prompt_logprobs,
            "outputs": outputs,
            "finished": finished,
            "metrics": seq_group.metrics,
            "lora_request": seq_group.lora_request,
            "encoder_prompt": encoder_prompt,
            "encoder_prompt_token_ids": encoder_prompt_token_ids,
            "num_cached_tokens": num_cached_tokens,
            "multi_modal_placeholders": seq_group.multi_modal_placeholders
        }
328
329
330

        if use_cache:
            request_output = seq_group.cached_request_output
331
            request_output.__init__(**init_kwargs)  # type: ignore
332
        else:
333
            request_output = cls(**init_kwargs)  # type: ignore
334
335

        return request_output
336
337
338
339
340

    def __repr__(self) -> str:
        return (f"RequestOutput(request_id={self.request_id}, "
                f"prompt={self.prompt!r}, "
                f"prompt_token_ids={self.prompt_token_ids}, "
341
342
                f"encoder_prompt={self.encoder_prompt!r}, "
                f"encoder_prompt_token_ids={self.encoder_prompt_token_ids}, "
343
                f"prompt_logprobs={self.prompt_logprobs}, "
344
                f"outputs={self.outputs}, "
345
                f"finished={self.finished}, "
346
                f"metrics={self.metrics}, "
347
                f"lora_request={self.lora_request}, "
348
349
                f"num_cached_tokens={self.num_cached_tokens}, "
                f"multi_modal_placeholders={self.multi_modal_placeholders})")
350
351


352
353
354
355
_O = TypeVar("_O", default=PoolingOutput)


class PoolingRequestOutput(Generic[_O]):
356
    """
357
    The output data of a pooling request to the LLM.
358
359

    Args:
360
361
        request_id (str): A unique identifier for the pooling request.
        outputs (PoolingOutput): The pooling results for the given input.
362
        prompt_token_ids (List[int]): A list of token IDs used in the prompt.
363
        finished (bool): A flag indicating whether the pooling is completed.
364
365
    """

366
    def __init__(self, request_id: str, outputs: _O,
367
368
369
370
371
372
                 prompt_token_ids: List[int], finished: bool):
        self.request_id = request_id
        self.prompt_token_ids = prompt_token_ids
        self.finished = finished
        self.outputs = outputs

373
374
375
376
377
    @staticmethod
    def from_seq_group(seq_group: SequenceGroup) -> "PoolingRequestOutput":
        pooled_data = seq_group.pooled_data
        assert pooled_data is not None

378
379
        data = pooled_data.to(dtype=torch.float32, device="cpu")
        output = PoolingOutput(data)
380
381
382
        prompt_token_ids = seq_group.prompt_token_ids
        finished = seq_group.is_finished()

383
384
        return PoolingRequestOutput(seq_group.request_id, output,
                                    prompt_token_ids, finished)
385
386
387

    def __repr__(self):
        """
388
        Returns a string representation of an PoolingRequestOutput instance.
389
390

        The representation includes the request_id and the number of outputs,
391
        providing a quick overview of the pooling request's results.
392
393

        Returns:
394
            str: A string representation of the PoolingRequestOutput instance.
395
        """
396
397
        return (f"{type(self).__name__}(request_id={self.request_id!r}, "
                f"outputs={self.outputs!r}, "
398
399
400
401
                f"prompt_token_ids={self.prompt_token_ids}, "
                f"finished={self.finished})")


402
403
404
405
406
407
408
409
410
411
412
413
414
class RequestOutputFactory:

    @staticmethod
    def create(seq_group: SequenceGroup,
               seq_id_to_seq_group: Dict[str, SequenceGroupBase],
               use_cache: bool = False):
        if seq_group.pooled_data is not None:
            return PoolingRequestOutput.from_seq_group(seq_group)
        else:
            return RequestOutput.from_seq_group(seq_group, use_cache,
                                                seq_id_to_seq_group)


415
@dataclass
416
417
class EmbeddingOutput:
    """The output data of one embedding output of a request.
418
419

    Args:
420
421
        embedding: The embedding vector, which is a list of floats.
        Its length depends on the hidden dimension of the model.
422
    """
423
424
425
426
427
428
429
430
431
432
433
434
435
    embedding: list[float]

    @staticmethod
    def from_base(pooling_output: PoolingOutput):
        pooled_data = pooling_output.data
        if pooled_data.ndim != 1:
            raise ValueError("pooled_data should be a 1-D embedding vector")

        return EmbeddingOutput(pooled_data.tolist())

    @property
    def hidden_size(self) -> int:
        return len(self.embedding)
436
437

    def __repr__(self) -> str:
438
        return f"EmbeddingOutput(hidden_size={self.hidden_size})"
439
440


441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
class EmbeddingRequestOutput(PoolingRequestOutput[EmbeddingOutput]):

    @staticmethod
    def from_base(request_output: PoolingRequestOutput):
        return EmbeddingRequestOutput(
            request_id=request_output.request_id,
            outputs=EmbeddingOutput.from_base(request_output.outputs),
            prompt_token_ids=request_output.prompt_token_ids,
            finished=request_output.finished,
        )


@dataclass
class ClassificationOutput:
    """The output data of one classification output of a request.
456
457

    Args:
458
459
        probs: The probability vector, which is a list of floats.
        Its length depends on the number of classes.
460
    """
461
    probs: list[float]
462

463
464
465
466
467
    @staticmethod
    def from_base(pooling_output: PoolingOutput):
        pooled_data = pooling_output.data
        if pooled_data.ndim != 1:
            raise ValueError("pooled_data should be a 1-D probability vector")
468

469
        return ClassificationOutput(pooled_data.tolist())
470

471
472
473
    @property
    def num_classes(self) -> int:
        return len(self.probs)
474

475
476
    def __repr__(self) -> str:
        return f"ClassificationOutput(num_classes={self.num_classes})"
477
478


479
class ClassificationRequestOutput(PoolingRequestOutput[ClassificationOutput]):
480
481

    @staticmethod
482
483
484
485
486
487
488
    def from_base(request_output: PoolingRequestOutput):
        return ClassificationRequestOutput(
            request_id=request_output.request_id,
            outputs=ClassificationOutput.from_base(request_output.outputs),
            prompt_token_ids=request_output.prompt_token_ids,
            finished=request_output.finished,
        )
489
490


491
492
493
@dataclass
class ScoringOutput:
    """The output data of one scoring output of a request.
494

495
496
497
498
499
500
501
502
503
504
    Args:
        score: The similarity score, which is a scalar value.
    """
    score: float

    @staticmethod
    def from_base(pooling_output: PoolingOutput):
        pooled_data = pooling_output.data
        if pooled_data.ndim != 0:
            raise ValueError("pooled_data should be a scalar score")
505

506
        return ScoringOutput(pooled_data.item())
507

508
509
    def __repr__(self) -> str:
        return f"ScoringOutput(score={self.score})"
510

511
    @property
512
513
    @deprecated("`LLM.score()` now returns scalar scores. "
                "Please access it via the `score` attribute. ")
514
515
    def embedding(self) -> list[float]:
        return [self.score]
516
517


518
519
520
521
522
523
524
525
526
527
class ScoringRequestOutput(PoolingRequestOutput[ScoringOutput]):

    @staticmethod
    def from_base(request_output: PoolingRequestOutput):
        return ScoringRequestOutput(
            request_id=request_output.request_id,
            outputs=ScoringOutput.from_base(request_output.outputs),
            prompt_token_ids=request_output.prompt_token_ids,
            finished=request_output.finished,
        )