[Feature]: PyTorch workflow does not support bad words

[tonyay163]

🚀 The feature, motivation and pitch

SamplingParams.bad_token_ids seems to work with SamplingType = TRTLLMSampler but not TorchSampler

Alternatives

TRTLLMSampler, but it also doesn't support everything TorchSampler, namely some of the newer speculative decoding configs

Additional context

No response

Before submitting a new issue...

• Make sure you already searched for relevant issues, and checked the documentation and examples for answers to frequently asked questions.

[tonyay163]

I think I can use embedding bias as a workaround since it's implemented, maybe like:

self.embedding_bias = torch.zeros(len(self.tokenizer), dtype=torch.float32)
for token_id in self.bad_token_ids:
    self.embedding_bias[token_id - 1] = float("-inf")
...

[karljang]

Thanks for reporting this!
This appears to be a missing piece from our transition to the PyTorch workflow. I've shared this with the team.
In the meantime, if you'd like to contribute a fix, we'd be happy to review it, it looks like you already have some ideas on how to implement it. Contributions are always welcome!

[github-actions]

Issue has not received an update in over 14 days. Adding stale label.

[github-actions]

This issue was closed because it has been 14 days without activity since it has been marked as stale.