File size: 3,882 Bytes
3b57fd1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 | #!/usr/bin/env python3
"""Generate text from this pack. Run: python quickstart.py "your prompt here"
This pack carries its own loader because the weights need transforms that stock
MLX loaders do not apply. That is why `mlx_lm.load()` reports an unsupported
model type; there is no patched mlx-lm to install.
Packs with a vision tower (config.json `components.vision`) load through
runtime/vision_artifact.py on mlx-vlm; text-only packs through runtime/artifact.py.
The sampler comes from generation_config.json and is passed explicitly, because
mlx-lm and mlx-vlm do not read that file.
"""
import json
import sys
from pathlib import Path
PACK = Path(__file__).resolve().parent
sys.path.insert(0, str(PACK / "runtime"))
MAX_TOKENS = 256
def read_json(name, default):
path = PACK / name
return json.loads(path.read_text()) if path.is_file() else default
def sampler_settings():
gen = read_json("generation_config.json", {})
out = {"temperature": gen.get("temperature", 1.0), "top_p": gen.get("top_p", 0.95), "top_k": gen.get("top_k", 20)}
if gen.get("do_sample") is False:
out["temperature"] = 0.0
if gen.get("min_p"):
out["min_p"] = gen["min_p"]
if gen.get("repetition_penalty") not in (None, 1.0):
out["repetition_penalty"] = gen["repetition_penalty"]
return out
def render(prompt):
template_path = PACK / "chat_template.jinja"
if not template_path.is_file():
return prompt
from jinja2.sandbox import ImmutableSandboxedEnvironment
return (
ImmutableSandboxedEnvironment(trim_blocks=True, lstrip_blocks=True)
.from_string(template_path.read_text())
.render(
messages=[{"role": "user", "content": prompt}],
add_generation_prompt=True,
enable_thinking=False,
)
)
def generate_vision(prompt, settings):
from mlx_vlm import generate
from vision_artifact import load_vl_model
model, processor, _ = load_vl_model(PACK)
result = generate(model, processor, render(prompt), max_tokens=MAX_TOKENS, **settings)
return result if isinstance(result, str) else result.text
def generate_text(prompt, settings):
import mlx.core as mx
from artifact import load_model
from mlx_lm.sample_utils import make_logits_processors, make_sampler
from tokenizers import Tokenizer
model, _ = load_model(PACK)
tokenizer = Tokenizer.from_file(str(PACK / "tokenizer.json"))
eos = read_json("generation_config.json", {}).get("eos_token_id")
stop = set(eos if isinstance(eos, list) else [] if eos is None else [eos])
stop |= {i for i in (tokenizer.token_to_id("<|im_end|>"), tokenizer.token_to_id("<|endoftext|>")) if i is not None}
sample = make_sampler(
temp=settings["temperature"], top_p=settings["top_p"], top_k=settings["top_k"],
min_p=settings.get("min_p", 0.0),
)
processors = make_logits_processors(repetition_penalty=settings.get("repetition_penalty"))
x = mx.array([tokenizer.encode(render(prompt), add_special_tokens=False).ids])
cache = model.make_cache()
generated = []
for _ in range(MAX_TOKENS):
logits = model.lm_head(model.model(x, cache=cache)[:, -1:, :])[:, -1, :]
for processor in processors:
logits = processor(mx.array(generated), logits)
x = sample(logits)[:, None]
mx.eval(x)
token = int(x.item())
if token in stop:
break
generated.append(token)
return tokenizer.decode(generated)
def main():
prompt = " ".join(sys.argv[1:]) or "What is 2 plus 2? Answer with the number only."
settings = sampler_settings()
if read_json("config.json", {}).get("components", {}).get("vision"):
print(generate_vision(prompt, settings))
else:
print(generate_text(prompt, settings))
if __name__ == "__main__":
main()
|