14 lines
425 B
Python
14 lines
425 B
Python
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)
|