Add beam search

[Brentably]

Hey, I'm considering opening a PR to add beam search.

I suppose I would try to follow huggingface/transformers's spec closely.

Would this type of thing be welcomed?

Also wanted to see if anyone had any thoughts about how to best do this.

[awni]

Would this type of thing be welcomed?

Could you say more about what you would use it for / what it would be useful for in general?

I think a beam search adds a lot of complexity and it's not entirely clear to me what it buys us. One way to prototype new features like that is to build it as a feature on top of mlx-lm. And then if it's useful and you still want to do so we can upstream the diff back here.

[Brentably]

Honestly most of my motivation is to learn more about the internals of libraries like this, and just become a better programmer. That and beam search is interesting to me for a particular inferencing usecase I've been working on, where I need to search for the most valid paths forward.

I digress. The most relevant thread here to look up is probably vllm's thread where they consider dropping support for beam search: vllm-project/vllm#6226

The TLDR is that it's still widely used in production systems and people still find it improves performance / reliability for a number of cases.

[awni]

Thanks for linking the thread! That's super interesting. It does validate my concern on the extra complexity of beam search. I

Honestly most of my motivation is to learn more about the internals of libraries like this, and just become a better programmer.

I would encourage you to implement it on top of mlx-lm. Should help achieve your goals. You could even upstream changes in the underlying models / APIs to make an implementation like that easier to do on top of mlx-lm.