Add logits_processor option to generate_step function

[nathanrchn]

This update introduces token masking capabilities to the generate_step function via a new logits_processor parameter. This enhancement supports constrained decoding scenarios that require token masking prior to sampling.

Updates include:

1. New logits_processor parameter in generate_step function
2. Token masking logic implemented within _step function
3. Updated docstring for generate_step to describe logits_processor
4. Created mx.array of all tokens including prompt

Usage example:

def logits_processor(input_ids: mx.array, logits: mx.array) -> mx.array:
    return grammar_processor(input_ids, logits)

Here, grammar_processor could represent a custom constrained decoding approach.

[awni]

Sorry for the delayed review. This is cool! And I think we can include it but I want to clarify something first. We have the logit_bias argument already which can do some of what the logit_processor does but in a less flexible way. Is that insufficient for your use cases? If so.. could you explain?

In addition, I don't think we need to support both arguments as it's a bit messy and redundant. Perhaps we can remove the logit_bias and keep the more flexible logit_processor if needed.

[nathanrchn]

I can indeed use the logit_bias argument but as my project primarily uses the transformers library, it would be nice to have the same logit_processor. Additionally, I would need to rewrite the stream_generate or generate method because the mask changes with each forward pass. Moreover, when masking tokens for constrained decoding, the majority of tokens are typically masked. Creating the logit_bias dictionary might slow down the generation process slightly, as it needs to cover thousands of tokens.

If you agree, I can remove the logit_bias from the arguments.

[awni]

Sounds good, thanks for clarifying. Let's remove logit_bias in favor of logit_processor then. Thanks!

[awni]

Is this a bug? Shouldn't it be tokens?

[nathanrchn]

Yes of course.

[awni]

Thanks for the addition!

[awni]

I added logit_bias back because it's part of the OpenAI API spec and I don't want to break compatability with that.

As a follow up, a nice thing to do would be to refactor out logit_bias and repition_penalty out of generate_step since they can both use the logit_processor. That would simplify generate_step quite nicely.

[nathanrchn]

Do you mean adding logit_bias and repition_penalty as argument for the generation and stream_generate and create logits_processor method to handle the biases and the repetition penalty?

If so, it might be beneficial to modify the logits_processor type from a single function to a list of functions to apply sequentially.

[awni]

I hadn't thought it through too carefully but yes what your describing is more or less what I had in mind.

If so, it might be beneficial to modify the logits_processor type from a single function to a list of functions to apply sequentially.

Yea .. that may be cleaner.