File size: 5,133 Bytes
ed99a2e | 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 | #!/usr/bin/env python3
"""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()
|