File size: 2,169 Bytes
4e12f7c 55bd028 4e12f7c 55bd028 4e12f7c 55bd028 4e12f7c 55bd028 4e12f7c 55bd028 4e12f7c | 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 | """Run a GLiNER2 Core ML classifier without loading the original model weights."""
import argparse
import json
from pathlib import Path
import coremltools as ct
import numpy as np
from preprocessing import load_processor, prepare_with_processor
def classify(
model_dir: str,
text: str,
task: str,
labels: list[str],
length: int = 128,
max_options: int = 8,
*,
package: str | None = None,
units: str = "all",
):
model_dir = Path(model_dir)
default_package = f"gliner2_small_classification_fp16_L{length}_K{max_options}.mlpackage"
package_path = Path(package or default_package)
if not package_path.is_absolute():
package_path = model_dir / package_path
compute_units = {"all": ct.ComputeUnit.ALL, "cpu-ane": ct.ComputeUnit.CPU_AND_NE}[units]
processor = load_processor(str(model_dir / "tokenizer"))
arrays = prepare_with_processor(processor, text, task, labels, length, max_options)
model = ct.models.MLModel(str(package_path), compute_units=compute_units)
scores = np.asarray(model.predict(arrays)["probabilities"])[0, : len(labels)]
return {"label": labels[int(scores.argmax())], "confidence": float(scores.max()),
"probabilities": {label: float(score) for label, score in zip(labels, scores)}}
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--model-dir", required=True)
parser.add_argument("--text", required=True)
parser.add_argument("--task", default="decision")
parser.add_argument("--labels", required=True, help="JSON list of label strings")
parser.add_argument("--length", type=int, default=128)
parser.add_argument("--max-options", type=int, default=8)
parser.add_argument("--package", help="Package name inside model-dir, or absolute package path")
parser.add_argument("--units", choices=["all", "cpu-ane"], default="all")
args = parser.parse_args()
print(json.dumps(classify(args.model_dir, args.text, args.task, json.loads(args.labels),
args.length, args.max_options, package=args.package, units=args.units), indent=2))
if __name__ == "__main__":
main()
|