From 548d670960f0f398d6031af9e9afb00409c29c97 Mon Sep 17 00:00:00 2001 From: Samuel Sloniker Date: Sat, 26 Nov 2022 17:50:26 -0800 Subject: [PATCH] Use Classifier for --category --- gptc/__main__.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/gptc/__main__.py b/gptc/__main__.py index 6c3e99b..11cd795 100644 --- a/gptc/__main__.py +++ b/gptc/__main__.py @@ -74,14 +74,12 @@ def main() -> None: else: text = sys.stdin.read() - probabilities = model.confidence(text, args.max_ngram_length) if args.category: - try: - print(sorted(probabilities.items(), key=lambda x: x[1])[-1][0]) - except IndexError: - print(None) + classifier = gptc.Classifier(model, args.max_ngram_length) + print(classifier.classify(text)) else: + probabilities = model.confidence(text, args.max_ngram_length) print(json.dumps(probabilities)) else: print(json.dumps(gptc.pack(args.model, True)[0]))