Error 'Received invalid kth 2along axis -1 for array with shape: (1,2)' when generate using mixtral model with MLX format 4 quantized. about mlx-examples HOT 7 CLOSED

huggingface-cli download --resume-download cloudyu/Yi-34Bx2-MoE-60B-DPO

from mlx_lm import convert
convert('cloudyu/Yi-34Bx2-MoE-60B-DPO', mlx_path='mlx-community/Yi-34Bx2-MoE-60B-DPO-4bit-mlx', quantize=True, q_bits=4)

# file test_hello.py
from mlx_lm import load, generate
import argparse
from transformers import AutoTokenizer, AutoModel, AutoModelForCausalLM

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="Convert models on huggingface hub to mlx format.")
    parser.add_argument("repo_id", help="repo_id on huggingface.co")
    parser.add_argument("-mlx", "--mlx", help="True if model is mlx format", action="store_true", default=False)
    parser.add_argument("-chat", "--chat", help="True if model is a chat model", action="store_true", default=False)
    parser.add_argument("--prompt", help="Prompt", default="hello")
    args = parser.parse_args()

    is_mlx = args.mlx
    repo_id = args.repo_id
    is_chat_model = args.chat
    prompt = args.prompt

    if is_mlx:
        model, tokenizer = load(repo_id)
        response = generate(model, tokenizer, prompt=prompt, verbose=True)
        print(response)
    else:
        if is_chat_model:
            tokenizer = AutoTokenizer.from_pretrained(repo_id, trust_remote_code=True)
            model = AutoModel.from_pretrained(repo_id, trust_remote_code=True).to('mps')
            model = model.eval()
            response, history = model.chat(tokenizer, prompt, history=[])
            print(response)
        else:
            tokenizer = AutoTokenizer.from_pretrained(repo_id)
            model = AutoModelForCausalLM.from_pretrained(repo_id)
            input_ids = tokenizer(prompt, return_tensors="pt")
            outputs = model.generate(**input_ids, max_new_tokens=10000)
            print(tokenizer.decode(outputs[0]))

==========
Prompt: hello
Traceback (most recent call last):
  File "/Users/hurongliang/git/llmresearch/test_hello.py", line 20, in <module>
    response = generate(model, tokenizer, prompt=prompt, verbose=True)
  File "/Users/hurongliang/miniconda3/envs/llm/lib/python3.10/site-packages/mlx_lm/utils.py", line 214, in generate
    for (token, prob), n in zip(
  File "/Users/hurongliang/miniconda3/envs/llm/lib/python3.10/site-packages/mlx_lm/utils.py", line 157, in generate_step
    logits, cache = model(y[None], cache=cache)
  File "/Users/hurongliang/miniconda3/envs/llm/lib/python3.10/site-packages/mlx_lm/models/mixtral.py", line 260, in __call__
    out, cache = self.model(inputs, cache)
  File "/Users/hurongliang/miniconda3/envs/llm/lib/python3.10/site-packages/mlx_lm/models/mixtral.py", line 243, in __call__
    h, cache[e] = layer(h, mask, cache[e])
  File "/Users/hurongliang/miniconda3/envs/llm/lib/python3.10/site-packages/mlx_lm/models/mixtral.py", line 209, in __call__
    r = self.block_sparse_moe(self.post_attention_layernorm(h))
  File "/Users/hurongliang/miniconda3/envs/llm/lib/python3.10/site-packages/mlx_lm/models/mixtral.py", line 160, in __call__
    mx.argpartition(-gates, kth=ne, axis=-1)[:, :ne]
ValueError: [argpartition] Received invalid kth 2along axis -1 for array with shape: (1,2)