blakeslee-research/nucleic-purpose-classifier
3
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge nucleic/sleek-ember-seal-uady into dev

Browse Source
This commit is contained in:
abkslm
2026-07-30 20:56:46 -07:00
parent e05521482f
commit 9e462a79fc
9 changed files with 832 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
eval.py
+130 -8
View File
@@ -13,6 +13,8 @@ from collections import Counter
from pathlib import Path
from typing import Any, Sequence
import numpy as np
from purpose_data import (
HARD_SLICES,
LABELS,
@@ -100,6 +102,16 @@ def _synchronize(torch: Any, device: Any) -> None:
torch.mps.synchronize()
def _coreml_compute_unit(coremltools: Any, requested: str) -> Any:
values = {
"all": coremltools.ComputeUnit.ALL,
"cpu-only": coremltools.ComputeUnit.CPU_ONLY,
"cpu-and-gpu": coremltools.ComputeUnit.CPU_AND_GPU,
"cpu-and-ne": coremltools.ComputeUnit.CPU_AND_NE,
}
return values[requested]
def routing_tier_drift(
records: Sequence[dict[str, Any]],
actual: Sequence[int],
@@ -245,8 +257,62 @@ def evaluate(args: argparse.Namespace) -> dict[str, Any]:
label_to_id = {label: index for index, label in enumerate(LABELS)}
tokenizer = AutoTokenizer.from_pretrained(args.model_dir, local_files_only=True)
onnx_session = None
coreml_model = None
reference_model = None
if args.onnx_model is not None:
reference_onnx_session = None
reference_onnx_input_names: set[str] = set()
if args.compare_pytorch and args.compare_onnx is not None:
raise DataError("choose only one parity reference")
if args.coreml_model is not None:
try:
import coremltools as ct
except ImportError as exc:
raise DataError(
"Core ML evaluation requires requirements-coreml.txt on macOS"
) from exc
if args.device != "auto":
raise DataError(
"Core ML compute placement uses --coreml-compute-units, not --device"
)
coreml_model = ct.models.MLModel(
str(args.coreml_model),
compute_units=_coreml_compute_unit(ct, args.coreml_compute_units),
)
coreml_input_names = {
item.name for item in coreml_model.get_spec().description.input
}
device = torch.device("cpu")
model = None
onnx_input_names = set()
runtime_name = f"coreml-{args.coreml_compute_units}"
if args.compare_pytorch:
reference_model = AutoModelForSequenceClassification.from_pretrained(
args.model_dir,
local_files_only=True,
).to(device)
reference_model.eval()
if args.compare_onnx is not None:
try:
import onnxruntime as ort
except ImportError as exc:
raise DataError(
"Core ML↔ONNX parity requires onnxruntime"
) from exc
reference_options = ort.SessionOptions()
reference_options.graph_optimization_level = (
ort.GraphOptimizationLevel.ORT_ENABLE_ALL
)
reference_onnx_session = ort.InferenceSession(
str(args.compare_onnx),
sess_options=reference_options,
providers=["CPUExecutionProvider"],
)
reference_onnx_input_names = {
item.name for item in reference_onnx_session.get_inputs()
}
elif args.onnx_model is not None:
if args.compare_onnx is not None:
raise DataError("--compare-onnx requires --coreml-model")
try:
import onnxruntime as ort
except ImportError as exc:
@@ -273,8 +339,12 @@ def evaluate(args: argparse.Namespace) -> dict[str, Any]:
).to(device)
reference_model.eval()
else:
if args.compare_onnx is not None:
raise DataError("--compare-onnx requires --coreml-model")
if args.compare_pytorch:
raise DataError("--compare-pytorch requires --onnx-model")
raise DataError(
"--compare-pytorch requires --onnx-model or --coreml-model"
)
device = _device(torch, args.device)
model = AutoModelForSequenceClassification.from_pretrained(
args.model_dir, local_files_only=True
@@ -284,6 +354,14 @@ def evaluate(args: argparse.Namespace) -> dict[str, Any]:
runtime_name = "pytorch"
def predict_logits(encoded: dict[str, Any]) -> Any:
if coreml_model is not None:
inputs = {
key: value.numpy().astype(np.int32, copy=False)
for key, value in encoded.items()
if key in coreml_input_names
}
output = np.asarray(coreml_model.predict(inputs)["logits"])
return torch.from_numpy(output.reshape(1, len(LABELS)))
if onnx_session is not None:
inputs = {
key: value.numpy()
@@ -300,7 +378,9 @@ def evaluate(args: argparse.Namespace) -> dict[str, Any]:
confidences: list[str] = []
margins: list[float] = []
reference_predictions: list[int] = []
inference_batch_size = 1 if onnx_session is not None else args.batch_size
inference_batch_size = (
1 if onnx_session is not None or coreml_model is not None else args.batch_size
)
with torch.inference_mode():
for start in range(0, len(records), inference_batch_size):
batch = records[start : start + inference_batch_size]
@@ -315,6 +395,19 @@ def evaluate(args: argparse.Namespace) -> dict[str, Any]:
reference_predictions.extend(
reference_logits.argmax(dim=-1).tolist()
)
elif reference_onnx_session is not None:
reference_inputs = {
key: value.numpy()
for key, value in encoded.items()
if key in reference_onnx_input_names
}
reference_logits = reference_onnx_session.run(
["logits"],
reference_inputs,
)[0]
reference_predictions.extend(
np.asarray(reference_logits).argmax(axis=-1).tolist()
)
distribution = torch.softmax(logits, dim=-1)
top = torch.topk(distribution, k=2, dim=-1)
batch_probabilities = top.values[:, 0].tolist()
@@ -465,7 +558,7 @@ def evaluate(args: argparse.Namespace) -> dict[str, Any]:
"modelVersion": calibration.get("modelVersion", args.model_dir.name),
"device": str(device),
"runtime": runtime_name,
"artifact": str(args.onnx_model or args.model_dir),
"artifact": str(args.coreml_model or args.onnx_model or args.model_dir),
"fixedInputShape": [1, MAX_LENGTH],
"overall": metrics,
"scoredClassification": scored_metrics,
@@ -508,13 +601,21 @@ def evaluate(args: argparse.Namespace) -> dict[str, Any]:
},
"misclassifications": misclassifications,
}
if reference_model is not None:
report["pytorchParity"] = prediction_agreement(
if reference_predictions:
parity_name = (
"onnxParity" if reference_onnx_session is not None else "pytorchParity"
)
parity = prediction_agreement(
records,
actual,
reference_predictions,
predicted,
)
report[parity_name] = parity
if reference_onnx_session is not None:
report["gates"]["onnxLabelAgreementAtLeast99_5Percent"] = (
parity["scoredLabelAgreement"] >= 0.995
)
write_json(args.report, report)
return report
@@ -522,15 +623,31 @@ def evaluate(args: argparse.Namespace) -> dict[str, Any]:
def build_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--model-dir", type=Path, default=DEFAULT_MODEL_DIR)
parser.add_argument(
runtime = parser.add_mutually_exclusive_group()
runtime.add_argument(
"--onnx-model",
type=Path,
help="score a fixed-shape ONNX artifact instead of the PyTorch checkpoint",
)
runtime.add_argument(
"--coreml-model",
type=Path,
help="score a fixed-shape Core ML package instead of the PyTorch checkpoint",
)
parser.add_argument(
"--compare-pytorch",
action="store_true",
help="include label-level drift from --model-dir when scoring ONNX",
help="include label-level drift from --model-dir when scoring an artifact",
)
parser.add_argument(
"--compare-onnx",
type=Path,
help="include Core ML label-level drift from this ONNX reference",
)
parser.add_argument(
"--coreml-compute-units",
choices=("all", "cpu-only", "cpu-and-gpu", "cpu-and-ne"),
default="cpu-and-ne",
)
parser.add_argument("--calibration", type=Path, default=DEFAULT_CALIBRATION)
parser.add_argument("--test", type=Path, default=DEFAULT_TEST)
@@ -552,6 +669,11 @@ def main(argv: Sequence[str] | None = None) -> int:
args = parser.parse_args(argv)
if args.batch_size <= 0 or args.latency_samples < 0:
parser.error("batch size must be positive and latency samples non-negative")
if args.coreml_model is not None and args.compare_onnx is None and not args.no_gate:
parser.error(
"gated Core ML evaluation requires --compare-onnx; use --no-gate only "
"for diagnostics"
)
try:
report = evaluate(args)
except (DataError, OSError, ValueError) as exc: