File size: 1,736 Bytes
811d51e | 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 | #!/usr/bin/env python3
import json
import sys
from pathlib import Path
import numpy as np
import onnxruntime as ort
REPO_ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(REPO_ROOT))
from features import build_features # noqa: E402
TASK_LABELS = [
"general_chat", "writing", "translation", "summarization", "research",
"coding", "mathematics", "document_analysis", "high_stakes",
]
COMPLEXITY_LABELS = ["low", "medium", "high"]
def distribution(labels, values):
return {label: float(value) for label, value in zip(labels, values, strict=True)}
def classify(text: str) -> dict:
model_path = REPO_ROOT / "model.onnx"
session = ort.InferenceSession(str(model_path), providers=["CPUExecutionProvider"])
features = build_features(text).astype(np.float32).reshape(1, -1)
outputs = session.run(None, {"features": features})
task_probs = outputs[0][0]
complexity_probs = outputs[1][0]
task_index = int(np.argmax(task_probs))
complexity_index = int(np.argmax(complexity_probs))
return {
"task": {
"label": TASK_LABELS[task_index],
"confidence": float(task_probs[task_index]),
"probabilities": distribution(TASK_LABELS, task_probs),
},
"complexity": {
"label": COMPLEXITY_LABELS[complexity_index],
"confidence": float(complexity_probs[complexity_index]),
"probabilities": distribution(COMPLEXITY_LABELS, complexity_probs),
},
}
def main() -> None:
if len(sys.argv) < 2:
raise SystemExit('Usage: python examples/python/classify.py "your prompt"')
print(json.dumps(classify(" ".join(sys.argv[1:])), indent=2))
if __name__ == "__main__":
main()
|