File size: 3,247 Bytes
628e2fd | 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 | """Run the GLiNER2.5-Decide 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_decision, task_labels
def package_name(precision: str, length: int, max_heads: int, max_options: int) -> str:
return f"gliner2_decide_classification_{precision}_L{length}_H{max_heads}_K{max_options}.mlpackage"
def decode(tasks: dict, logits: np.ndarray) -> dict:
"""Apply the native activation and label rules to per-head logits of shape (heads, options)."""
results = {}
for head, (name, labels) in enumerate(task_labels(tasks).items()):
config = tasks[name] if isinstance(tasks[name], dict) else {}
values = logits[head, : len(labels)].astype(np.float64)
activation = config.get("class_act", "auto")
multi = config.get("multi_label", False)
if activation == "sigmoid" or (activation == "auto" and multi):
probs = 1.0 / (1.0 + np.exp(-values))
else:
probs = np.exp(values - values.max())
probs /= probs.sum()
if multi:
threshold = config.get("cls_threshold", 0.5)
chosen = [{"label": labels[j], "confidence": float(probs[j])} for j in range(len(labels))
if probs[j] >= threshold]
best = int(probs.argmax())
results[name] = chosen or [{"label": labels[best], "confidence": float(probs[best])}]
else:
best = int(probs.argmax())
results[name] = {"label": labels[best], "confidence": float(probs[best])}
return results
class CoreMLDecide:
def __init__(self, model_dir: str, precision: str = "fp16", length: int = 128, max_heads: int = 4,
max_options: int | None = None, compute_units=ct.ComputeUnit.ALL):
model_dir = Path(model_dir)
# Published buckets: L128 holds 8 labels per head, L256 and L512 hold 32.
max_options = max_options or (8 if length == 128 else 32)
self.length, self.max_heads, self.max_options = length, max_heads, max_options
self.processor = load_processor(str(model_dir))
package = model_dir / package_name(precision, length, max_heads, max_options)
self.model = ct.models.MLModel(str(package), compute_units=compute_units)
def classify(self, text: str, tasks: dict) -> dict:
arrays = prepare_decision(self.processor, text, tasks, self.length, self.max_heads, self.max_options)
logits = np.asarray(self.model.predict(arrays)["logits"])[0]
return decode(tasks, logits)
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--model-dir", required=True)
parser.add_argument("--text", required=True)
parser.add_argument("--tasks", required=True, help='JSON object, e.g. {"intent": ["a", "b"]}')
parser.add_argument("--precision", default="fp16")
parser.add_argument("--length", type=int, default=128)
args = parser.parse_args()
model = CoreMLDecide(args.model_dir, args.precision, args.length)
print(json.dumps(model.classify(args.text, json.loads(args.tasks)), indent=2))
if __name__ == "__main__":
main()
|