System instructions & pre-fill

[fblissjr]

Been thinking a bit about an elegant way to allow for modifying system messages / system instructions in the generate functions for instruct/chat models, as well as pre-fill - essentially creating a message chain of: System Role -> User Role -> Assistant Role -> n. This is often a very underused but powerful way to work with single step instruct models (we do this often with Anthropic's pre-fill features in the API)

Here's an example of both:

Adding only --system-message arg:

# Example use case with formatting rules
mlx_lm.generate \
  --model mlx-community/Llama-3.2-3B-Instruct-bf16 \
  --system-message "You are a data extraction expert. Always format your response as a JSON object with fields: 'company_name', 'revenue', 'year'. Only include factual information mentioned in the text. Use null for missing fields." \
  --prompt "According to reports, Apple's revenue reached $383.93 billion during their fiscal year 2023."

# Output:
{
  "company_name": "Apple",
  "revenue": 383930000000,
  "year": 2023
}

Adding messages-file for pre-fill:

// messages.json
{
  "messages": [
    {
      "role": "system",
      "content": "You are a data extraction expert. Format your response as a JSON object with fields: 'company_name', 'revenue', 'year'. Use null for missing fields."
    },
    {
      "role": "user",
      "content": "Microsoft generated $211.9 billion in revenue in 2022."
    },
    {
      "role": "assistant",
      "content": "{\"company_name\": \"Microsoft\", \"revenue\": 211900000000, \"year\": 2022}"
    },
    {
      "role": "user", 
      "content": "The gaming division of Sony had strong performance in Q3."
    },
    {
      "role": "assistant",
      "content": "{\"company_name\": \"Sony\", \"revenue\": null, \"year\": null}"
    },
    {
      "role": "user",
      "content": "According to reports, Apple's revenue reached $383.93 billion during their fiscal year 2023."
    }
  ]
}

mlx_lm.generate --messages-file messages.json

# Output is now more reliable due to examples / in context learning:
{
  "company_name": "Apple",
  "revenue": 383930000000,
  "year": 2023
}

For system role, I've been doing this for quite some time in my local repo for mlx-lm, and think this is probably the simplest approach:

1. Add argument to generate.py:

    parser.add_argument(
        "--system-message",
        type=str,
        help="System instructions for the model",
    )

2. Modify the messages building to support it:

if not args.ignore_chat_template and (
    hasattr(tokenizer, "apply_chat_template")
    and tokenizer.chat_template is not None
):
    messages = []
    if args.system_message:
        messages.append({"role": "system", "content": args.system_message})
    messages.append({
        "role": "user",
        "content": sys.stdin.read() if args.prompt == "-" else args.prompt,
    })
    prompt = tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )

However, to support the multi-turn pre-fill - - the best I can come up with is a messages JSON file, which could also work fantastically with the cache capabilities. Here's my initial proposed approach, but am open to simpler ones if they exist:

# messages file support - adds just one more argument
parser.add_argument(
    "--messages-file",
    type=str,
    help="JSON file containing full message history",
)

def main():
    messages = None
    if args.messages_file:
        with open(args.messages_file) as f:
            data = json.load(f)
            messages = data["messages"]
    else:
        messages = []
        if args.system_message:
            messages.append({"role": "system", "content": args.system_message})
        messages.append({
            "role": "user",
            "content": sys.stdin.read() if args.prompt == "-" else args.prompt
        })

Welcome any thoughts here, or if I should just open a PR to move discussion there?

[madroidmaq]

@fblissjr Perhaps you can try it. https://github.com/madroidmaq/mlx-omni-server