Axion-0.5B-MLX-LM

4 Bit Quantization

Py Script to try the Model

import os
from mlx_lm import generate, load

# Set this to either a Hugging Face repo ID or a folder containing a model.
MODEL_PATH_OR_REPO = "Rajiv714/Axion-0.5B-MLX-LM"

MAX_TOKENS = 1500

SYSTEM_PROMPT = "You are a helpful medical AI assistant."
USER_PROMPT = "I am having a fever for the past 10 days. What should I do?"

def get_model_path(model_path_or_repo):
    if os.path.isdir(model_path_or_repo):
        print(f"Using local model: {model_path_or_repo}")
        return model_path_or_repo

    print(f"Using Hugging Face model: {model_path_or_repo}")
    return model_path_or_repo

def main():
    model_path = get_model_path(MODEL_PATH_OR_REPO)

    print(f"Loading model: {model_path}\n")

    model, tokenizer = load(model_path)

    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": USER_PROMPT},
    ]

    prompt = tokenizer.apply_chat_template(messages, add_generation_prompt=True)

    generate(model, tokenizer, prompt=prompt, max_tokens=MAX_TOKENS, verbose=True)

if __name__ == "__main__":
    main()