hub_predict.py

#!/usr/bin/env python3

import argparse
import json

from hub_runtime import HubTicketClassifier


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("text")
    parser.add_argument("--repo-id")
    parser.add_argument("--revision", default="main")
    parser.add_argument("--top-k", type=int, default=3)
    args = parser.parse_args()

    classifier = HubTicketClassifier(
        repo_id=args.repo_id,
        revision=args.revision,
    )
    print(json.dumps(
        classifier.predict(args.text, args.top_k),
        ensure_ascii=False,
        indent=2,
    ))


if __name__ == "__main__":
    main()