File size: 1,953 Bytes
4dc5bc6 | 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 | #!/usr/bin/env python3
from __future__ import annotations
import argparse
import json
import os
from shieldstral_policy import (
discover_model,
flatten_taxonomy,
load_taxonomy,
node_result,
query_for,
score_policy,
)
def main() -> None:
parser = argparse.ArgumentParser(description="Score one Shieldstral hierarchy policy node.")
parser.add_argument("--node", required=True, help="Official SC/CAT ID or AXONVERTEX SUB ID")
parser.add_argument("--document", required=True)
parser.add_argument("--document-type", choices=("prompt", "response"), default="prompt")
parser.add_argument("--instruct", default="Evaluate whether the document matches the query criteria")
parser.add_argument("--threshold", type=float, default=0.5)
parser.add_argument("--base-url", default=os.getenv("BASE_URL", "http://127.0.0.1:18190/v1"))
parser.add_argument("--model", default=None)
parser.add_argument("--taxonomy", default=None)
parser.add_argument("--raw", action="store_true")
args = parser.parse_args()
taxonomy = load_taxonomy(args.taxonomy)
nodes = flatten_taxonomy(taxonomy)
node_id = args.node.upper()
if node_id not in nodes:
raise SystemExit(f"Unknown node ID: {node_id}")
node = nodes[node_id]
model = discover_model(args.base_url, args.model)
score, raw = score_policy(
base_url=args.base_url,
model=model,
instruct=args.instruct,
query=query_for(node, args.document_type),
document=args.document,
threshold=args.threshold,
)
result = {
"classification_mode": "single_policy_node",
"model": model,
"document_type": args.document_type,
"node": node_result(node, score, document_type=args.document_type),
}
if args.raw:
result["raw_response"] = raw
print(json.dumps(result, indent=2, ensure_ascii=False))
if __name__ == "__main__":
main()
|