49 lines
2.8 KiB
Python
49 lines
2.8 KiB
Python
"""Pinned MTPLX oracle for issue 107's real long two-image product request."""
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import subprocess
|
|
import sys
|
|
|
|
checkout, root, image, output = map(Path, sys.argv[1:5])
|
|
assert subprocess.check_output(["git", "-C", str(checkout), "rev-parse", "HEAD"], text=True).strip() == "9882703f3105363ddc37eca9f97aa09a1d387112"
|
|
sys.path.insert(0, str(checkout))
|
|
from mtplx.server.openai import _server_runtime_env_overrides
|
|
os.environ.update(_server_runtime_env_overrides(argparse.Namespace(model=str(root), generation_mode="mtp"), None))
|
|
os.environ["MTPLX_SUSTAINED_PREFILL"] = "1"
|
|
os.environ["MTPLX_PREFILL_CHUNK_SIZE"] = "2048"
|
|
import mlx.core as mx
|
|
import numpy as np
|
|
from mlx_lm.utils import load_model
|
|
from mtplx.models import qwen4_exp as qwen
|
|
from mtplx.runtime import _load_tokenizer_resilient, MTPLXRuntime
|
|
from mtplx.mtp_patch import MTPContract
|
|
from mtplx.sampling import SamplerConfig
|
|
from mtplx.generation import generate_mtpk
|
|
from mtplx.server.openai import ChatMessage, _encode_messages_uncached, _materialize_vision_splice
|
|
|
|
output.mkdir(parents=True, exist_ok=True)
|
|
print(json.dumps(dict(case="environment", overrides={k:v for k,v in os.environ.items() if k.startswith("MTPLX_")})), flush=True)
|
|
tokenizer = _load_tokenizer_resilient(root, json.loads((root/"config.json").read_text()))
|
|
model, _ = load_model(root, lazy=False, strict=True, get_model_classes=lambda **_: (qwen.Model,qwen.ModelArgs))
|
|
model.post_weight_load(root)
|
|
assert model.attach_mtp(root)
|
|
runtime = MTPLXRuntime(model, tokenizer, root, True, MTPContract())
|
|
text = "Describe these images" + "<|vision_start|><|image_pad|><|vision_end|>" * 2
|
|
ids = _encode_messages_uncached(tokenizer, [ChatMessage(role="user", content=text)], enable_thinking=True,
|
|
reasoning_effort="low", preserve_reasoning_history=True, tools=None)
|
|
ids, splice = _materialize_vision_splice(argparse.Namespace(args=argparse.Namespace(model=str(root))), [image.read_bytes()]*2, ids)
|
|
assert len(ids) > 2048
|
|
(output/"input.json").write_text(json.dumps(ids))
|
|
out = generate_mtpk(runtime, ids, max_tokens=1, sampler=SamplerConfig(temperature=0., top_p=.95, top_k=20),
|
|
speculative_depth=3, seed=1, stop_token_ids=set(tokenizer.eos_token_ids), mtp_history_policy="committed",
|
|
verify_strategy="batched", vision_splice=splice, capture_final_state=True,
|
|
prefill_callback=lambda event: print(json.dumps(dict(case="prefill", **event)), flush=True))
|
|
assert out.final_state.safe_to_commit
|
|
logits = np.asarray(out.final_state.final_logits.astype(mx.float32)).reshape(-1)
|
|
(output/"final.f32").write_bytes(logits.tobytes())
|
|
print(json.dumps(dict(case="long_vision", prompt_tokens=len(ids), tokens=out.tokens, text=out.text,
|
|
final_logits_sha256=hashlib.sha256(logits.tobytes()).hexdigest())), flush=True)
|