Output Logits

[Dookas]

I can see how to generate text with these models but I want to get the actual output logins over the vocab. Is there a simple way to do this at each step, do I need to create a custom detokenizer?

[davidkoski]

I think you want something like this:

https://github.com/ml-explore/mlx-swift-examples/blob/main/Libraries/MLXLMCommon/Evaluate.swift#L12

You can supply your own logit sampler (returning a single value MLXArray. If that isn't quite what you are looking for, then perhaps you want a generator of logits?

https://github.com/ml-explore/mlx-swift-examples/blob/main/Libraries/MLXLMCommon/Evaluate.swift#L244

https://github.com/ml-explore/mlx-swift-examples/blob/main/Libraries/MLXLMCommon/Evaluate.swift#L367

It would be easy to copy that and make a Sequence<MLXArray>. You need to be careful about evaluating those on the same thread they are produced in (synchronous with the iteration).

If this is something that is generally useful, perhaps the TokenIterator could take some kind of output preparer and we would have TokenIterator<R> with next() -> R?

[Dookas]

yeah the general idea is I want to enforce what the model can output. In my case I am trying to enforce only yes or no answers. I can do this by feeding the input and then looking at the first output logits. I can then respond with either yes or no depending on which token has higher probability assuming the vocab has a yes and no token.

[davidkoski]

OK, the logit sampler ought to be able to do that -- it has all the logits and its only task is to return a single token.