| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from collections.abc import Callable |
| 4 | from typing import TypeAlias |
| 5 | |
| 6 | import torch |
| 7 | |
| 8 | LogitsProcessor: TypeAlias = ( |
| 9 | Callable[[list[int], torch.Tensor], torch.Tensor] |
| 10 | | Callable[[list[int], list[int], torch.Tensor], torch.Tensor] |
| 11 | ) |
| 12 | """LogitsProcessor is a function that takes a list |
| 13 | of previously generated tokens, the logits tensor |
| 14 | for the next token and, optionally, prompt tokens as a |
| 15 | first argument, and returns a modified tensor of logits |
| 16 | to sample from.""" |