MLX: rodar LLM nativo em M3/M4 Apple Silicon

# MLX core — NumPy-like com lazy evaluation
import mlx.core as mx
import mlx.nn as nn

# Arrays não são materializados até mx.eval()
a = mx.random.normal(shape=(1024, 1024))
b = mx.random.normal(shape=(1024, 1024))
c = a @ b                       # lazy: grafo construído
d = mx.exp(c) + a               # lazy: fusão potencial
# Nenhum cálculo ainda

mx.eval(d)                      # materializa — agora computa
print(d.shape)                  # (1024, 1024)

# Diferenciação automática
def loss_fn(W, x, y):
    return mx.mean((x @ W - y) ** 2)

grad_fn = mx.grad(loss_fn)
W = mx.random.normal((10, 1))
x = mx.random.normal((100, 10))
y = mx.random.normal((100, 1))

dW = grad_fn(W, x, y)           # gradiente vs primeiro arg

# Composability: gradients of gradients
hessian_diag = mx.grad(mx.grad(loss_fn))(W, x, y)

# Instalação
pip install mlx-lm

# Download e geração — modelo já quantizado em INT4
python -m mlx_lm.generate \
    --model mlx-community/Meta-Llama-3.1-70B-Instruct-4bit \
    --prompt "Explique PagedAttention em 3 bullets." \
    --max-tokens 500 \
    --temp 0.5

# Conversão própria de HF → MLX INT4
python -m mlx_lm.convert \
    --hf-path meta-llama/Meta-Llama-3.1-8B-Instruct \
    --mlx-path ./Llama-3.1-8B-Instruct-mlx-q4 \
    --quantize \
    -q --q-bits 4 --q-group-size 64

# Servidor compatível com OpenAI API
python -m mlx_lm.server \
    --model mlx-community/Qwen2.5-72B-Instruct-4bit \
    --port 8080 \
    --temp 0.5

# curl http://localhost:8080/v1/chat/completions \
#   -d '{"model":"mlx","messages":[{"role":"user","content":"oi"}]}'

# API Python idiomática
from mlx_lm import load, generate, stream_generate

model, tokenizer = load("mlx-community/Meta-Llama-3.1-8B-Instruct-4bit")

prompt = tokenizer.apply_chat_template(
    [{"role": "user", "content": "O que é unified memory em M3 Ultra?"}],
    tokenize=False,
    add_generation_prompt=True,
)

# Generation simples
response = generate(model, tokenizer, prompt=prompt, max_tokens=500, temp=0.5)
print(response)

# Streaming
for tok in stream_generate(model, tokenizer, prompt=prompt, max_tokens=500):
    print(tok.text, end="", flush=True)

# Speculative decoding nativo
from mlx_lm import generate
draft_model, _ = load("mlx-community/Llama-3.2-1B-Instruct-4bit")

response = generate(
    model, tokenizer, prompt=prompt, max_tokens=500,
    draft_model=draft_model, num_draft_tokens=5,
)

# mlx_lm.lora — QLoRA-style em hardware Apple
python -m mlx_lm.lora \
    --model mlx-community/Meta-Llama-3.1-70B-Instruct-4bit \
    --train \
    --data ./data_jsonl \          # JSONL com {"prompt":..., "completion":...}
    --iters 1000 \
    --batch-size 2 \
    --lora-layers 16 \             # quantos layers receber LoRA (top N)
    --learning-rate 1e-5 \
    --adapter-path ./ffv-tutor-lora \
    --save-every 200

# Fundir LoRA com base
python -m mlx_lm.fuse \
    --model mlx-community/Meta-Llama-3.1-70B-Instruct-4bit \
    --adapter-path ./ffv-tutor-lora \
    --save-path ./ffv-tutor-70b-fused-mlx

# Servir o modelo fundido
python -m mlx_lm.server --model ./ffv-tutor-70b-fused-mlx

# Modelos multimodais (VLMs) no MLX
from mlx_vlm import load, generate

model, processor = load("mlx-community/Qwen2.5-VL-7B-Instruct-4bit")

response = generate(
    model, processor,
    image="https://example.com/diagram.png",
    prompt="Descreva o diagrama de arquitetura.",
    max_tokens=500,
    temp=0.3,
)

# Suportados:
# - LLaVA family (1.5, 1.6, NeXT)
# - Qwen2-VL, Qwen2.5-VL
# - Llama 3.2 Vision (11B, 90B)
# - Phi-3-vision, Phi-3.5-vision
# - InternVL2
# - Idefics3