Implementation of recommend

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 '<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,39 +55,17 @@ 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
-                })
-    
-    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}")
-
+            filename = os.path.basename(path)
+            print_recommendation(path, filename, ctx)
     
 
 @ucsinfer.command('gather')