Skip to content

vllm.v1.sample.ops.penalties

_convert_to_tensors 

_convert_to_tensors(
    output_token_ids: list[list[int]],
    vocab_size: int,
    device: device,
) -> Tensor

Convert the different list data structures to tensors.

Source code in vllm/v1/sample/ops/penalties.py 
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
def _convert_to_tensors(output_token_ids: list[list[int]], vocab_size: int,
                        device: torch.device) -> torch.Tensor:
    """
    Convert the different list data structures to tensors.
    """
    output_tokens_tensor = make_tensor_with_pad(
        output_token_ids,
        # Use the value of vocab_size as a pad since we don't have a
        # token_id of this value.
        pad=vocab_size,
        device="cpu",
        dtype=torch.int64,
        pin_memory=is_pin_memory_available(),
    )
    return output_tokens_tensor.to(device, non_blocking=True)

apply_all_penalties 

apply_all_penalties(
    logits: Tensor,
    prompt_token_ids: Tensor,
    presence_penalties: Tensor,
    frequency_penalties: Tensor,
    repetition_penalties: Tensor,
    output_token_ids: list[list[int]],
) -> Tensor

Applies presence, frequency and repetition penalties to the logits.

Source code in vllm/v1/sample/ops/penalties.py 
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
def apply_all_penalties(
    logits: torch.Tensor,
    prompt_token_ids: torch.Tensor,
    presence_penalties: torch.Tensor,
    frequency_penalties: torch.Tensor,
    repetition_penalties: torch.Tensor,
    output_token_ids: list[list[int]],
) -> torch.Tensor:
    """
    Applies presence, frequency and repetition penalties to the logits.
    """
    _, vocab_size = logits.shape
    output_tokens_t = _convert_to_tensors(output_token_ids, vocab_size,
                                          logits.device)
    return apply_penalties(logits, prompt_token_ids, output_tokens_t,
                           presence_penalties, frequency_penalties,
                           repetition_penalties)