Implementation of recommend

This commit is contained in:
Jamie Hardt
2025-08-27 15:37:53 -07:00
parent de7ad3d65a
commit 739d27fe71

View File

@@ -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 '<None>'}")
for r in recommendations:
cat, subcat, _ = ctx.lookup_category(r)
print(f"- {r}: {cat}-{subcat}")
@click.group(epilog="For more information see "
"<https://git.squad51.us/jamie/ucsinfer>")
# @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 '<None>'}")
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'),