import argparse

from mini_lm import predict


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Run token classification inference.")
    parser.add_argument("text", nargs="*", help="Input text to label.")
    parser.add_argument("--top-k", type=int, default=3)
    args = parser.parse_args()

    text = " ".join(args.text).strip() or "convert foo.mp4 to mkv"
    predict(text, top_k=args.top_k)