biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[Core]: Option To Use Prompt Token Ids Inside Logits Processor (#4985)

Browse Source 
Co-authored-by: Elisei Smirnov <el.smirnov@innopolis.university>
This commit is contained in:
Elisei Smirnov
2024-05-24 01:04:24 +03:00
committed by GitHub
parent a1242324c9
commit e3470f8753
2 changed files with 24 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
vllm/model_executor/layers/logits_processor.py
View File

@@ -1,4 +1,5 @@
"""A layer that compute logits from hidden_stats."""
import inspect
from typing import Optional
import torch
@@ -95,15 +96,25 @@ def _apply_logits_processors(
seq_ids = seq_group.seq_ids
sampling_params = seq_group.sampling_params
logits_processors = sampling_params.logits_processors
if logits_processors:
found_logits_processors = True
for seq_id, logits_row_idx in zip(seq_ids,
seq_group.sample_indices):
logits_row = logits[logits_row_idx]
token_ids = seq_group.seq_data[seq_id].output_token_ids
past_tokens_ids = seq_group.seq_data[seq_id].output_token_ids
prompt_tokens_ids = seq_group.seq_data[seq_id].prompt_token_ids
for logits_processor in logits_processors:
logits_row = logits_processor(token_ids, logits_row)
parameters = inspect.signature(logits_processor).parameters
if len(parameters) == 3:
logits_row = logits_processor(prompt_tokens_ids,
past_tokens_ids,
logits_row)
else:
logits_row = logits_processor(past_tokens_ids,
logits_row)
logits[logits_row_idx] = logits_row
logits_processed += len(seq_group.sample_indices) + len(