Starcoder2 Support

[lazarust]

It would be nice if we could support starcoder2.

Expected Behavior:

python -m mlx_lm.generate --model bigcode/starcoder2-3b --prompt "Write a quick sort in C++"

[awni]

Happy to accept a contribution for that!

[lazarust]

@awni I'll take a stab at it and see what I can do! Could you point me to a good starting point?

[awni]

Of course!

So it's relatively simple, you add the model file to this directory: https://github.com/ml-explore/mlx-examples/tree/main/llms/mlx_lm/models

You can see other examples in there, I recommend starting from one of the models that is pretty similar to starcoder2

Make sure the name of the model file is the same as the model_type in the config (starcoder2) (https://huggingface.co/bigcode/starcoder2-7b/blob/main/config.json#L17)

If you want to add LoRA support you can edit this section of the code as needed

[Vaibhavs10]

Wohoo! btw - Just a heads-up the model is a FIM type. So would be a little different in terms of the way you prompt it.

In case you want to see the changes from Mistral to StarCoder2: huggingface/transformers#29120

[awni]

Just a heads-up the model is a FIM type

What's FIM? And does the HF tokenizer template handle that or do we need to do something special?

[lazarust]

@awni Could you also assign me to this issue? I can't seem to do it

[Vaibhavs10]

What's FIM? And does the HF tokenizer template handle that or do we need to do something special?

FIM - Fill in the Middle. Nothing special. Just that it is not chat-tuned.

https://huggingface.co/bigcode/starcoder2-15b#running-the-model-on-cpugpumulti-gpu

[sugatoray]

FIM: Fill in the Middle arxiv-paper