I want use the function prefix_allowed_tokens_fn of huggingface model.generate(), where of vllm's source code shall I modify?

[zoubaihan]

Hello, we all know that in huggingface transformers' origin model.generate() method, we can set the function paremeter prefix_allowed_tokens_fn to restrict the generation rule. I want to use this function in vllm just like I used in origin model.generate() to control the generation process, could you please tell me where of the source code shall I modify to make the model generation obey my custom prefix_allowed_tokens_fn?

[zhuohan123]

This is related to #379. You can add logit_bias as a parameter of SamplingParams. Then you can add this bias when sampling for the results to not generate some specific tokens, similar to what we are doing for embedding_bias from the models:

vllm/vllm/model_executor/layers/sampler.py

Lines 48 to 49 in 20b0d88

	if embedding_bias is not None:
	logits += embedding_bias

Feel free to contribute to this feature! It will be very helpful for many use cases.

[Etelis]

I will try implementing this!

Thanks.

[hmellor]

Functionality of prefix_allowed_tokens_fn can be achieved using logit processes, as per #2434 (comment)