PC: Add support for Starcoder2

[Blaizzy]

Hey @awni and @Muhtasham,

I was super excited about integrating starcoder 2 into mlx. However, while working on the mlx-langchain integration, I realized that @Muhtasham beat me to the punch and submitted a PR. Kudos to you!

I took a look at the original draft PR a few days ago and noticed a few areas that could use some improvement (mlp and others). When I tested the code with the quantized models uploaded to the hub, it didn't quite work as expected. But not to worry, I've fixed it! ✨ ✅

This implementation is working well and matches the original model using transformers in full precision and quantized configurations.

python -m mlx_lm.generate --model mlx-community/starcoder2-3b-4bit --prompt "Write a quick sort in C++"  --colorize --max-tokens 200

Output:

Prompt: Write a quick sort in C++

```cpp
//https://www.hackerrank.com/challenges/quicksort3/problem

#include <bits/stdc++.h>

using namespace std;

vector<int> quickSort(vector<int> ar) {
    vector<int> output;
    int pivot = ar[0];
    for(int i = 1; i < ar.size(); i++){
        if(ar[i] < pivot)
            output.push_back(ar[i]);
    }
    output.push_back(pivot);
    for(int i = 1; i < ar.size(); i++){
        if(ar[i] > pivot)
            output.push_back(ar[i]);
    }
    return output;
}

int main(void) {
    vector <int>  ar;
    int num_of_elements;
    cin >> num

[mzbac]

Nice work! Just curious, does this work with the FIM prompt?

[mzbac]

Maybe we need to keep scores.astype(mx.float32) in case softmax overflows on float16 values.

[awni]

@Blaizzy thanks for adding! I just merged #502 since it was basically done. I'm going to close this in favor of that. Appreciate the insights here which helped with #502!

[Blaizzy]

Nice work! Just curious, does this work with the FIM prompt?

Thanks! It should work without a problem.

If not I can investigate.

[Blaizzy]

@Blaizzy thanks for adding! I just merged #502 since it was basically done. I'm going to close this in favor of that. Appreciate the insights here which helped with #502!

Most welcome!

It's my pleasure 😊.