How to get the logits in mlx-lm?

[Mihaiii]

I'm interested in adding mlx-lm as a provider in Backtrack Sampler.

I generate one token at a time, apply my own temperature, and handle the selection of the next token in my own code. Therefore, I need to retrieve the logits from mlx-lm rather than pass it a custom sampler.

I have some changes in a branch, fwiw, but they currently don't work because I handle log probabilities (logprobs) instead of logits, as I'm using the utils.stream_generate function. I don't necessarily need to use stream_generate; I used it because I saw it here.

Is there a way to get the logits in mlx-lm?

[awni]

Typically it should be possible to treat the logprobs as if they were logits. Unless you are doing something that relies on the normalization term (which is not so common).

For example for temperature scaling, you should be able to use the logprobs as if they were logits since it should be the case that:mx.softmax(temp * logits) == mx.softmax(temp * logprobs).

Another option is to send stream_generate a custom logits_processor using the logits_processors keyword arg.

It would be good to try and track down where the discrepancy in your code is. If it's actually the case that you need to operate on the logits as opposed the log probs.. would be interesting to know what the reason is.