| |
| """Offline native MLX text completion for OpenSML-Base-175M.""" |
| from pathlib import Path |
| import argparse |
| import hashlib |
| import json |
| import math |
| import sys |
| import tempfile |
| import time |
| import zipfile |
| ROOT = Path(__file__).resolve().parent |
| BUNDLE_SHA256 = "14f85f0ff677d35a6fb34f118f2dd21a3fdd504fff27c163860ab919b21e0e98" |
|
|
| def sha256(path): |
| with Path(path).open("rb") as file: |
| return hashlib.file_digest(file, "sha256").hexdigest() |
|
|
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| inputs = parser.add_mutually_exclusive_group(required=True) |
| inputs.add_argument("--prompt", help="Raw text prefix; no chat template is applied") |
| inputs.add_argument("--prompt-file", type=Path, help="UTF-8 text prefix file") |
| parser.add_argument("--weights", type=Path, default=ROOT/"model.safetensors") |
| parser.add_argument("--max-tokens", type=int, default=128) |
| parser.add_argument("--temperature", type=float, default=0) |
| parser.add_argument("--seed", type=int, default=42) |
| parser.add_argument("--json", action="store_true", help="Print completion and generation metadata as JSON") |
| args = parser.parse_args() |
| if not 1 <= args.max_tokens <= 2047: |
| parser.error("--max-tokens must be between 1 and 2047") |
| if not math.isfinite(args.temperature) or args.temperature < 0: |
| parser.error("--temperature must be finite and nonnegative") |
| prompt = args.prompt if args.prompt is not None else args.prompt_file.read_text(encoding="utf-8") |
| config = json.loads((ROOT/"config.json").read_text()) |
| if not args.weights.is_file(): |
| parser.error("Missing model.safetensors; download the complete repository") |
| if sha256(args.weights) != config["weight_sha256"]: |
| parser.error("Weights do not match the published base model identity") |
| archive = ROOT/"native_bundle.zip" |
| if sha256(archive) != BUNDLE_SHA256: |
| parser.error("Native source bundle failed its integrity check") |
| with tempfile.TemporaryDirectory(prefix="opensml-base-runtime-") as directory: |
| runtime = Path(directory) |
| with zipfile.ZipFile(archive) as bundle: |
| for member in bundle.infolist(): |
| p = Path(member.filename) |
| if p.is_absolute() or ".." in p.parts or (member.external_attr >> 16) & 0o170000 == 0o120000: |
| parser.error("Unsafe source-bundle entry") |
| bundle.extractall(runtime) |
| sys.path.insert(0, str(runtime/"pretraining")) |
| from sml_v1.tokenization import Tokenizer |
| from sml_v1.model import TransformerConfig, TransformerLM, count_parameters |
| import mlx.core as mx |
| from mlx.utils import tree_map |
| tokenizer = Tokenizer(runtime/"tokenizer") |
| ids = tokenizer.encode(prompt) |
| if not ids: |
| parser.error("Prompt must contain at least one token") |
| if len(ids)+args.max_tokens > config["context_limit"]: |
| parser.error("Prompt plus generation budget exceeds 2048 tokens; no truncation is performed") |
| mx.set_default_device(mx.gpu) |
| mx.set_cache_limit(4*1024**3) |
| mx.random.seed(args.seed) |
| model = TransformerLM(TransformerConfig(**config["model"])) |
| model.update(tree_map(lambda value:value.astype(mx.bfloat16),model.parameters())) |
| model.load_weights(str(args.weights)) |
| mx.eval(model.parameters()) |
| model.eval() |
| if count_parameters(model) != config["parameters"]: |
| raise ValueError("Unexpected parameter count") |
| inputs = mx.array([ids],dtype=mx.int32) |
| cache = None |
| generated = [] |
| stop = "token_limit" |
| start = time.perf_counter() |
| try: |
| for _ in range(args.max_tokens): |
| logits, cache = model.logits(inputs, cache) |
| final = logits[0,-1] |
| if args.temperature == 0: |
| token = int(mx.argmax(final).item()) |
| else: |
| token = int(mx.random.categorical(final.astype(mx.float32)/args.temperature).item()) |
| if token == tokenizer.eos: |
| stop = "document_end"; break |
| if token in (0,2,3): |
| stop = "structural_token"; break |
| generated.append(token) |
| inputs = mx.array([[token]],dtype=mx.int32) |
| completion = tokenizer.decode(generated) |
| record = {"model":config["model_name"],"prompt":prompt,"completion":completion, |
| "prompt_tokens":len(ids),"output_tokens":len(generated), |
| "stop_reason":stop,"seconds":time.perf_counter()-start, |
| "protocol":{"serialization":"raw text prefix","temperature":args.temperature, |
| "seed":args.seed,"max_tokens":args.max_tokens,"context_limit":2048, |
| "truncation":False,"compute_dtype":"bfloat16"}} |
| print(json.dumps(record,indent=2,ensure_ascii=False) if args.json else completion) |
| finally: |
| mx.clear_cache() |
|
|
| if __name__ == "__main__": |
| main() |
|
|
