355 lines
12 KiB
Python
355 lines
12 KiB
Python
#!/usr/bin/env python3
|
|
"""Convert the selected purpose-lite checkpoint to a fixed-shape Core ML package."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
import math
|
|
import shutil
|
|
import sys
|
|
from pathlib import Path
|
|
from typing import Any, Sequence
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn.functional as functional
|
|
|
|
from purpose_data import LABELS, DataError, load_jsonl, write_json
|
|
from train import MAX_LENGTH, encode_fixed_shape
|
|
|
|
|
|
SCRIPT_DIR = Path(__file__).resolve().parent
|
|
DEFAULT_MODEL_DIR = (
|
|
SCRIPT_DIR
|
|
/ "outputs"
|
|
/ "purpose-lite-v1-distilled-qat-mlx-4e"
|
|
/ "model"
|
|
)
|
|
DEFAULT_VALIDATION = SCRIPT_DIR / ".artifacts" / "dataset-v1" / "validation.jsonl"
|
|
DEFAULT_OUTPUT = (
|
|
SCRIPT_DIR
|
|
/ "outputs"
|
|
/ "purpose-lite-v1-distilled-qat-mlx-4e"
|
|
/ "coreml"
|
|
/ "purpose-lite-v1-fp16.mlpackage"
|
|
)
|
|
COREMLTOOLS_VERSION = "9.0"
|
|
|
|
|
|
class FixedShapeBertForCoreML(torch.nn.Module):
|
|
"""Conversion-only BERT forward without Transformers' dynamic mask helpers."""
|
|
|
|
def __init__(self, model: Any) -> None:
|
|
super().__init__()
|
|
self.model = model
|
|
config = model.config
|
|
self.num_heads = int(config.num_attention_heads)
|
|
self.head_size = int(config.hidden_size) // self.num_heads
|
|
if int(config.hidden_size) % self.num_heads:
|
|
raise DataError("BERT hidden size must be divisible by attention heads")
|
|
if int(config.max_position_embeddings) < MAX_LENGTH:
|
|
raise DataError("BERT checkpoint cannot represent the fixed 128-token input")
|
|
self.register_buffer(
|
|
"fixed_position_ids",
|
|
torch.arange(MAX_LENGTH, dtype=torch.int64).reshape(1, MAX_LENGTH),
|
|
persistent=False,
|
|
)
|
|
|
|
def forward(
|
|
self,
|
|
input_ids: Any,
|
|
attention_mask: Any,
|
|
token_type_ids: Any,
|
|
) -> Any:
|
|
bert = self.model.bert
|
|
input_ids = input_ids.to(torch.int64)
|
|
token_type_ids = token_type_ids.to(torch.int64)
|
|
value = (
|
|
bert.embeddings.word_embeddings(input_ids)
|
|
+ bert.embeddings.position_embeddings(self.fixed_position_ids)
|
|
+ bert.embeddings.token_type_embeddings(token_type_ids)
|
|
)
|
|
value = bert.embeddings.LayerNorm(value)
|
|
zero = torch.zeros((), dtype=value.dtype, device=value.device)
|
|
hidden = torch.full((), -10000.0, dtype=value.dtype, device=value.device)
|
|
additive_mask = torch.where(
|
|
attention_mask[:, None, None, :] != 0,
|
|
zero,
|
|
hidden,
|
|
)
|
|
|
|
for layer in bert.encoder.layer:
|
|
self_attention = layer.attention.self
|
|
|
|
def split_heads(projected: Any) -> Any:
|
|
return projected.reshape(
|
|
1,
|
|
MAX_LENGTH,
|
|
self.num_heads,
|
|
self.head_size,
|
|
).permute(0, 2, 1, 3)
|
|
|
|
queries = split_heads(self_attention.query(value))
|
|
keys = split_heads(self_attention.key(value))
|
|
values = split_heads(self_attention.value(value))
|
|
scores = torch.matmul(queries, keys.transpose(-1, -2)) / math.sqrt(
|
|
self.head_size
|
|
)
|
|
probabilities = torch.softmax(scores + additive_mask, dim=-1)
|
|
context = torch.matmul(probabilities, values)
|
|
context = context.permute(0, 2, 1, 3).reshape(
|
|
1,
|
|
MAX_LENGTH,
|
|
-1,
|
|
)
|
|
attention_output = layer.attention.output.dense(context)
|
|
value = layer.attention.output.LayerNorm(value + attention_output)
|
|
intermediate = functional.gelu(
|
|
layer.intermediate.dense(value),
|
|
approximate="none",
|
|
)
|
|
value = layer.output.LayerNorm(
|
|
value + layer.output.dense(intermediate)
|
|
)
|
|
|
|
pooled = torch.tanh(bert.pooler.dense(value[:, 0]))
|
|
return self.model.classifier(pooled)
|
|
|
|
|
|
def _checkpoint_config(model_dir: Path) -> dict[str, Any]:
|
|
config_path = model_dir / "config.json"
|
|
try:
|
|
config = json.loads(config_path.read_text(encoding="utf-8"))
|
|
except (OSError, UnicodeError, json.JSONDecodeError) as exc:
|
|
raise DataError(f"{config_path}: cannot load model config: {exc}") from exc
|
|
configured_labels = [
|
|
config.get("id2label", {}).get(
|
|
str(index),
|
|
config.get("id2label", {}).get(index),
|
|
)
|
|
for index in range(len(LABELS))
|
|
]
|
|
if (
|
|
config.get("model_type") != "bert"
|
|
or config.get("hidden_size") != 384
|
|
or config.get("num_hidden_layers") != 6
|
|
):
|
|
raise DataError("Core ML conversion requires the purpose-lite BERT architecture")
|
|
if configured_labels != list(LABELS):
|
|
raise DataError("Core ML checkpoint label order does not match purpose-lite")
|
|
return config
|
|
|
|
|
|
def _sha256(path: Path) -> str:
|
|
digest = hashlib.sha256()
|
|
with path.open("rb") as handle:
|
|
for chunk in iter(lambda: handle.read(1024 * 1024), b""):
|
|
digest.update(chunk)
|
|
return digest.hexdigest()
|
|
|
|
|
|
def _package_manifest(package: Path) -> dict[str, Any]:
|
|
files = []
|
|
for path in sorted(item for item in package.rglob("*") if item.is_file()):
|
|
files.append(
|
|
{
|
|
"path": str(path.relative_to(package)),
|
|
"bytes": path.stat().st_size,
|
|
"sha256": _sha256(path),
|
|
}
|
|
)
|
|
return {
|
|
"schemaVersion": 1,
|
|
"modelVersion": "purpose-lite-v1",
|
|
"package": package.name,
|
|
"bytes": sum(item["bytes"] for item in files),
|
|
"files": files,
|
|
}
|
|
|
|
|
|
def _verify_wrapper_parity(
|
|
torch_model: Any,
|
|
wrapper: FixedShapeBertForCoreML,
|
|
tokenizer: Any,
|
|
records: Sequence[dict[str, Any]],
|
|
) -> float:
|
|
encoded = encode_fixed_shape(
|
|
tokenizer,
|
|
[record["prompt"] for record in records],
|
|
torch,
|
|
)
|
|
maximum_error = 0.0
|
|
with torch.inference_mode():
|
|
for index in range(len(records)):
|
|
item = {key: value[index : index + 1] for key, value in encoded.items()}
|
|
reference = torch_model(**item).logits
|
|
candidate = wrapper(
|
|
item["input_ids"].to(torch.int32),
|
|
item["attention_mask"].to(torch.int32),
|
|
item.get("token_type_ids", torch.zeros_like(item["input_ids"])).to(
|
|
torch.int32
|
|
),
|
|
)
|
|
error = float(torch.max(torch.abs(reference - candidate)).item())
|
|
maximum_error = max(maximum_error, error)
|
|
torch.testing.assert_close(candidate, reference, rtol=1e-5, atol=2e-5)
|
|
if int(candidate.argmax(dim=-1).item()) != int(
|
|
reference.argmax(dim=-1).item()
|
|
):
|
|
raise DataError("conversion wrapper changed a parity prediction")
|
|
return maximum_error
|
|
|
|
|
|
def convert(args: argparse.Namespace) -> dict[str, Any]:
|
|
_checkpoint_config(args.model_dir)
|
|
checkpoint = args.model_dir / "model.safetensors"
|
|
if not checkpoint.is_file():
|
|
raise DataError(f"{checkpoint}: checkpoint is missing")
|
|
if args.output.suffix != ".mlpackage":
|
|
raise DataError("Core ML output must end in .mlpackage")
|
|
if args.output.exists():
|
|
if not args.overwrite_output:
|
|
raise DataError(
|
|
f"{args.output}: output exists; pass --overwrite-output intentionally"
|
|
)
|
|
if args.output.is_dir():
|
|
shutil.rmtree(args.output)
|
|
else:
|
|
args.output.unlink()
|
|
args.output.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
try:
|
|
import coremltools as ct
|
|
from transformers import AutoModelForSequenceClassification, AutoTokenizer
|
|
except ImportError as exc:
|
|
raise DataError(
|
|
"Core ML conversion requires requirements-coreml.txt on macOS"
|
|
) from exc
|
|
if ct.__version__ != COREMLTOOLS_VERSION:
|
|
raise DataError(
|
|
f"expected coremltools {COREMLTOOLS_VERSION}, found {ct.__version__}"
|
|
)
|
|
|
|
model = AutoModelForSequenceClassification.from_pretrained(
|
|
args.model_dir,
|
|
local_files_only=True,
|
|
).eval()
|
|
tokenizer = AutoTokenizer.from_pretrained(args.model_dir, local_files_only=True)
|
|
wrapper = FixedShapeBertForCoreML(model).eval()
|
|
validation = load_jsonl(args.validation)
|
|
if len(validation) < args.parity_records:
|
|
raise DataError("validation split is smaller than --parity-records")
|
|
maximum_error = _verify_wrapper_parity(
|
|
model,
|
|
wrapper,
|
|
tokenizer,
|
|
validation[: args.parity_records],
|
|
)
|
|
print(f"conversion-wrapper parity: max_abs_error={maximum_error:.3g}")
|
|
|
|
example = (
|
|
torch.zeros((1, MAX_LENGTH), dtype=torch.int32),
|
|
torch.ones((1, MAX_LENGTH), dtype=torch.int32),
|
|
torch.zeros((1, MAX_LENGTH), dtype=torch.int32),
|
|
)
|
|
with torch.inference_mode():
|
|
traced = torch.jit.trace(wrapper, example, strict=True)
|
|
traced = torch.jit.freeze(traced)
|
|
deployment_target = getattr(ct.target, args.minimum_deployment_target, None)
|
|
if deployment_target is None:
|
|
raise DataError(
|
|
f"coremltools does not support {args.minimum_deployment_target}"
|
|
)
|
|
coreml_model = ct.convert(
|
|
traced,
|
|
convert_to="mlprogram",
|
|
minimum_deployment_target=deployment_target,
|
|
compute_precision=ct.precision.FLOAT16,
|
|
inputs=[
|
|
ct.TensorType(
|
|
name="input_ids",
|
|
shape=(1, MAX_LENGTH),
|
|
dtype=np.int32,
|
|
),
|
|
ct.TensorType(
|
|
name="attention_mask",
|
|
shape=(1, MAX_LENGTH),
|
|
dtype=np.int32,
|
|
),
|
|
ct.TensorType(
|
|
name="token_type_ids",
|
|
shape=(1, MAX_LENGTH),
|
|
dtype=np.int32,
|
|
),
|
|
],
|
|
outputs=[ct.TensorType(name="logits", dtype=np.float32)],
|
|
)
|
|
coreml_model.author = "Nucleic"
|
|
coreml_model.short_description = "purpose-lite-v1 prompt classifier"
|
|
coreml_model.version = "purpose-lite-v1"
|
|
coreml_model.user_defined_metadata["com.nucleic.model.version"] = (
|
|
"purpose-lite-v1"
|
|
)
|
|
coreml_model.user_defined_metadata["com.nucleic.model.labels"] = json.dumps(
|
|
list(LABELS),
|
|
separators=(",", ":"),
|
|
)
|
|
coreml_model.user_defined_metadata["com.nucleic.model.sourceSha256"] = _sha256(
|
|
checkpoint
|
|
)
|
|
coreml_model.user_defined_metadata["com.nucleic.model.fixedShape"] = "1x128"
|
|
coreml_model.user_defined_metadata["com.nucleic.model.minimumDeploymentTarget"] = (
|
|
args.minimum_deployment_target
|
|
)
|
|
coreml_model.save(str(args.output))
|
|
|
|
manifest = _package_manifest(args.output)
|
|
manifest.update(
|
|
{
|
|
"sourceCheckpoint": str(checkpoint),
|
|
"sourceCheckpointSha256": _sha256(checkpoint),
|
|
"coremltoolsVersion": ct.__version__,
|
|
"minimumDeploymentTarget": args.minimum_deployment_target,
|
|
"computePrecision": "float16",
|
|
"wrapperMaximumAbsoluteError": maximum_error,
|
|
}
|
|
)
|
|
manifest_path = args.output.with_name(f"{args.output.stem}-manifest.json")
|
|
write_json(manifest_path, manifest)
|
|
print(f"Core ML package: {args.output} ({manifest['bytes']} bytes)")
|
|
return manifest
|
|
|
|
|
|
def build_parser() -> argparse.ArgumentParser:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--model-dir", type=Path, default=DEFAULT_MODEL_DIR)
|
|
parser.add_argument("--validation", type=Path, default=DEFAULT_VALIDATION)
|
|
parser.add_argument("--output", type=Path, default=DEFAULT_OUTPUT)
|
|
parser.add_argument("--parity-records", type=int, default=16)
|
|
parser.add_argument(
|
|
"--minimum-deployment-target",
|
|
default="macOS15",
|
|
choices=("macOS15", "macOS26"),
|
|
)
|
|
parser.add_argument("--overwrite-output", action="store_true")
|
|
return parser
|
|
|
|
|
|
def main(argv: Sequence[str] | None = None) -> int:
|
|
parser = build_parser()
|
|
args = parser.parse_args(argv)
|
|
if args.parity_records <= 0:
|
|
parser.error("--parity-records must be positive")
|
|
try:
|
|
convert(args)
|
|
except (AssertionError, DataError, OSError, RuntimeError, ValueError) as exc:
|
|
print(f"error: {exc}", file=sys.stderr)
|
|
return 1
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|