37 lines
1.9 KiB
Python
37 lines
1.9 KiB
Python
"""Model-free #105 receipt from the pinned MTPLX/MLX oracle, never production."""
|
|
import json
|
|
import subprocess
|
|
from pathlib import Path
|
|
|
|
import mlx.core as mx
|
|
import numpy as np
|
|
from mtplx import generation
|
|
from mtplx.sampling import NonFiniteLogitsError, SamplerConfig
|
|
|
|
reference = Path(generation.__file__).resolve().parent.parent
|
|
assert subprocess.check_output(["git", "-C", str(reference), "rev-parse", "HEAD"], text=True).strip() == "9882703f3105363ddc37eca9f97aa09a1d387112"
|
|
assert mx.__version__ == "0.32.2"
|
|
rows = {
|
|
"nan": np.full(64, np.nan, np.float32),
|
|
"positive_inf": np.array([np.inf] + [0.] * 63, np.float32),
|
|
"all_negative_inf": np.full(64, -np.inf, np.float32),
|
|
"valid_mask": np.array([-np.inf, 2., -np.inf, 1.] + [-np.inf] * 60, np.float32),
|
|
}
|
|
for name, values in rows.items():
|
|
row = mx.array(values)
|
|
for temperature in (0., .6, 1.):
|
|
config = SamplerConfig(temperature=temperature, top_p=.95, top_k=20)
|
|
try:
|
|
token, _ = generation._sample_from_logits(row, config, np.random.default_rng(0))
|
|
result = {"token": token}
|
|
except NonFiniteLogitsError as error:
|
|
result = {"error": str(error), "nan": error.nan_count, "inf": error.inf_count, "vocab": error.vocab_size}
|
|
print(json.dumps(dict(case="ordinary", row=name, temperature=temperature, **result)), flush=True)
|
|
config = SamplerConfig(temperature=1., top_p=.95, top_k=20)
|
|
token = int(generation._mx_lazy_sample(row, config, mx.random.key(7)).item())
|
|
print(json.dumps(dict(case="lazy", row=name, token=token)), flush=True)
|
|
# The specialized MTP greedy helper still lacks the ordinary guard at this
|
|
# revision. Keep the conflict with #105 visible; never label this parity.
|
|
token, _ = generation._greedy_draft_token_and_top_values(row.reshape(1,1,-1), topk=2)
|
|
print(json.dumps(dict(case="mtp_greedy_specialized", row=name, token=token)), flush=True)
|