93 lines
4.1 KiB
Python
93 lines
4.1 KiB
Python
"""Model-free #103 oracle receipts. Set PYTHONPATH to the immutable MTPLX checkout.
|
||
|
||
Run in the existing MLX 0.32.2 reference environment, passing the unchanged
|
||
Bare-Speed artifact folder. Only tokenizer assets are read; no model is loaded.
|
||
"""
|
||
import hashlib
|
||
import importlib.metadata
|
||
import json
|
||
import math
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
import mlx.core as mx
|
||
import numpy as np
|
||
from mtplx.generation import _mx_lazy_sample, _sample_from_logits
|
||
from mtplx.sampling import SamplerConfig, distribution_from_logits
|
||
from mtplx.runtime import _load_tokenizer_resilient
|
||
|
||
|
||
def emit(**receipt):
|
||
print(json.dumps(receipt, ensure_ascii=False), flush=True)
|
||
|
||
|
||
root = Path(sys.argv[1])
|
||
emit(case="environment", mlx=importlib.metadata.version("mlx"),
|
||
device=mx.device_info(),
|
||
hashes={n: hashlib.sha256((root/n).read_bytes()).hexdigest()
|
||
for n in ["config.json", "tokenizer.json", "merges.txt", "chat_template.jinja"]})
|
||
config = SamplerConfig(temperature=1., top_p=0.5, top_k=2)
|
||
values = np.array([math.log(.4), math.log(.3)] + [math.log(.3/62)]*62, dtype=np.float32)
|
||
p = distribution_from_logits(values.astype(np.float64), config)
|
||
emit(case="sampling_law", probabilities=p[:2].tolist())
|
||
counts = [0]*64
|
||
for seed in range(128):
|
||
token = _mx_lazy_sample(mx.array(values), config, mx.random.key(seed))
|
||
mx.eval(token)
|
||
counts[int(token.item())] += 1
|
||
emit(case="sampling_law_lazy", counts=counts[:2])
|
||
for temperature, top_k, top_p in [(.6,20,.95),(.2,8,.8),(1.,2,1.)]:
|
||
row = values * np.float32(temperature)
|
||
config = SamplerConfig(temperature=temperature, top_k=top_k, top_p=top_p)
|
||
probabilities = distribution_from_logits(row.astype(np.float64), config)
|
||
counts = [0]*64
|
||
for seed in range(512):
|
||
token = _mx_lazy_sample(mx.array(row), config, mx.random.key(seed))
|
||
mx.eval(token)
|
||
counts[int(token.item())] += 1
|
||
emit(case="sampling_matrix", temperature=temperature, top_k=top_k, top_p=top_p,
|
||
ordinary_probabilities=probabilities.tolist(), lazy_counts=counts)
|
||
for name, values in [
|
||
("nan", [float("nan")]*64),
|
||
("positive_inf", [float("inf")]+[0.]*63),
|
||
("all_negative_inf", [float("-inf")]*64),
|
||
("valid_mask", [float("-inf"),2.,float("-inf"),1.]+[float("-inf")]*60),
|
||
]:
|
||
for temperature in [0.,.6,1.]:
|
||
config = SamplerConfig(temperature=temperature, top_p=.95, top_k=20)
|
||
try:
|
||
token, _ = _sample_from_logits(mx.array(values), config, np.random.default_rng(0))
|
||
emit(case="invalid_logits", row=name, temperature=temperature, token=int(token))
|
||
except (ValueError, RuntimeError) as error:
|
||
emit(case="invalid_logits", row=name, temperature=temperature, error=type(error).__name__)
|
||
token = _mx_lazy_sample(mx.array(values), SamplerConfig(temperature=1., top_k=20, top_p=.95), mx.random.key(7))
|
||
mx.eval(token)
|
||
emit(case="invalid_logits_lazy", row=name, token=int(token.item()))
|
||
|
||
tokenizer = _load_tokenizer_resilient(root, json.loads((root/"config.json").read_text()))
|
||
for text in [
|
||
"भारत एक विशाल देश है जिसकी संस्कृति बहुत पुरानी है।",
|
||
"ประเทศไทยมีวัฒนธรรมที่เก่าแก่และสวยงามมาก",
|
||
"إِنَّ اللُّغَةَ الْعَرَبِيَّةَ جَمِيلَةٌ",
|
||
"The quick brown fox jumps over the lazy dog near the river bank.",
|
||
"def foo(x):\n return x + 1\n",
|
||
"中国的文化历史悠久。",
|
||
"a🙂b ©c ∑d —e …f €g ™h\n",
|
||
"x\u0301y १२३ ١٢٣ ①Ⅷ²\u200dZ",
|
||
" ①", " Ⅷ", " ²", " 𝟘",
|
||
]:
|
||
emit(case="tokenizer", text=text, ids=tokenizer.encode(text, add_special_tokens=False))
|
||
|
||
try:
|
||
from mtplx.attention_context import attention_phase
|
||
from mtplx.attention_math import attention_gate
|
||
except ImportError:
|
||
emit(case="bf16_verify_gate", available=False)
|
||
sys.exit(0)
|
||
gate = mx.array([-6.84375]*4, dtype=mx.bfloat16)
|
||
value = mx.ones((4,), dtype=mx.bfloat16)
|
||
with attention_phase("decode_verify"):
|
||
output = attention_gate(value, gate)
|
||
mx.eval(output)
|
||
emit(case="bf16_verify_gate", values=output.astype(mx.float32).tolist())
|