From 739d27fe711b4a68b0dd1d4809021672d7f2ca72 Mon Sep 17 00:00:00 2001 From: Jamie Hardt Date: Wed, 27 Aug 2025 15:37:53 -0700 Subject: [PATCH] Implementation of recommend --- ucsinfer/__main__.py | 44 +++++++++++++++++--------------------------- 1 file changed, 17 insertions(+), 27 deletions(-) diff --git a/ucsinfer/__main__.py b/ucsinfer/__main__.py index ed71dea..dad6ade 100644 --- a/ucsinfer/__main__.py +++ b/ucsinfer/__main__.py @@ -16,6 +16,18 @@ from .util import ffmpeg_description, parse_ucs def recommend_text(text: str, ctx: InferenceContext): return ctx.classify_text_ranked(text) +def print_recommendation(path: str | None, text: str, ctx: InferenceContext): + recommendations = ctx.classify_text_ranked(text) + print("----------") + if path: + print(f"Path: {path}") + + print(f"Text: {text or ''}") + for r in recommendations: + cat, subcat, _ = ctx.lookup_category(r) + print(f"- {r}: {cat}-{subcat}") + + @click.group(epilog="For more information see " "") # @click.option('--verbose', flag_value='verbose', help='Verbose output') @@ -43,40 +55,18 @@ def recommend(text, paths, model): m = SentenceTransformer(model) ctx = InferenceContext(m, model) - recommendations = [] if text is not None: - recommendations.append({ - "text": text, - "recommendations": recommend_text(text, ctx) - }) + print_recommendation(None, text, ctx) for path in paths: text = ffmpeg_description(path) if text: - recommendations.append({ - "path":path, - "text":text, - "recommendations":recommend_text(text, ctx) - }) + print_recommendation(path, text, ctx) + else: - recommendations.append({ - "path":path, - "text":None, - "recommendations":None - }) + filename = os.path.basename(path) + print_recommendation(path, filename, ctx) - for rec in recommendations: - print("----------") - if 'path' in rec: - print(f"Path: {rec['path']}") - - print(f"Text: {rec['text'] or ''}") - if rec['recommendations']: - for r in rec['recommendations']: - cat, subcat, _ = ctx.lookup_category(r) - print(f"- {r}: {cat}-{subcat}") - - @ucsinfer.command('gather') @click.option('--outfile', type=click.File(mode='w', encoding='utf8'),