diff --git a/README.md b/README.md index 46f63ce..611e6cc 100644 --- a/README.md +++ b/README.md @@ -8,21 +8,23 @@ GPTC provides both a CLI tool and a Python library. ### Classifying text - python -m gptc + python -m gptc classify + +This will prompt for a string and classify it, then print (in JSON) a dict of +the format `{category: probability, category:probability, ...}` to stdout. + +Alternatively, if you only need the most likely category, you can use this: + + python -m gptc classify [-c|--category] This will prompt for a string and classify it, outputting the category on stdout (or "None" if it cannot determine anything). -Alternatively, if you need confidence data, use: - - python -m gptc -j - -This will print (in JSON) a dict of the format `{category: probability, -category:probability, ...}` to stdout. - ### Compiling models - python -m gptc -c|--compile + python -m gptc compile + +This will print the compiled model in JSON to stdout. ## Library diff --git a/gptc/__main__.py b/gptc/__main__.py index d9627b5..350ae36 100644 --- a/gptc/__main__.py +++ b/gptc/__main__.py @@ -6,34 +6,44 @@ import json import sys import gptc -parser = argparse.ArgumentParser(description="General Purpose Text Classifier") -parser.add_argument("model", help="model to use") -parser.add_argument( - "-c", - "--compile", - help="compile raw model model to outfile", - metavar="outfile", -) -parser.add_argument( +parser = argparse.ArgumentParser(description="General Purpose Text Classifier", prog='gptc') +subparsers = parser.add_subparsers(dest="subparser_name", required=True) + +compile_parser = subparsers.add_parser('compile', help='compile a raw model') +compile_parser.add_argument("model", help="raw model to compile") + +classify_parser = subparsers.add_parser('classify', help='classify text') +classify_parser.add_argument("model", help="compiled model to use") +group = classify_parser.add_mutually_exclusive_group() +group.add_argument( "-j", - "--confidence", - help="output confidence dict in json", + "--json", + help="output confidence dict as JSON (default)", action="store_true", ) +group.add_argument( + "-c", + "--category", + help="output most likely category or `None`", + action="store_true", +) + args = parser.parse_args() with open(args.model, "r") as f: - raw_model = json.load(f) -if args.compile: - with open(args.compile, "w+") as f: - json.dump(gptc.compile(raw_model), f) + model = json.load(f) + +if args.subparser_name == 'compile': + print(json.dumps(gptc.compile(model))) else: - classifier = gptc.Classifier(raw_model) + classifier = gptc.Classifier(model) + if sys.stdin.isatty(): text = input("Text to analyse: ") else: text = sys.stdin.read() - if args.confidence: - print(json.dumps(classifier.confidence(text))) - else: + + if args.category: print(classifier.classify(text)) + else: + print(json.dumps(classifier.confidence(text)))