From 1a42d7d9f1c33fbe6668c67349a16c5a68495d5d Mon Sep 17 00:00:00 2001 From: Jamie Hardt Date: Sat, 30 Aug 2025 23:29:45 -0700 Subject: [PATCH] Refactored recommendation code into a new file --- ucsinfer/__main__.py | 61 +----------------------------------------- ucsinfer/recommend.py | 62 +++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 63 insertions(+), 60 deletions(-) create mode 100644 ucsinfer/recommend.py diff --git a/ucsinfer/__main__.py b/ucsinfer/__main__.py index 1e9b6b1..345b7c4 100644 --- a/ucsinfer/__main__.py +++ b/ucsinfer/__main__.py @@ -1,7 +1,6 @@ import os import sys import csv -from re import match from sentence_transformers import SentenceTransformer import tqdm @@ -9,67 +8,9 @@ import click from tabulate import tabulate, SEPARATING_LINE from .inference import InferenceContext, load_ucs +from .recommend import print_recommendation 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, - interactive_rename: bool): - """ - Print recommendations interactively. - - :returns: If interactive_rename is false or path is None, returns None. - If interactive, returns a tuple of bool, str | None, str | None where: - - if retval[0] is False, the user has requested processing quit. - - if retval[1] is a str, this is the text the user has entered to - perform a new inference instead of the file metadata. - `print_recommendation` should be called again with this argument. - - if retval[2] is a str, this is the catid the user has selected. - """ - recs = ctx.classify_text_ranked(text) - print("----------") - if path: - print(f"Path: {path}") - - print(f"Text: {text or ''}") - for i, r in enumerate(recs): - cat, subcat, _ = ctx.lookup_category(r) - print(f"- {i}: {r} ({cat}-{subcat})") - - if interactive_rename and path is not None: - response = input("#, t [text], ?, q > ") - - if m := match(r'^([0-9]+)', response): - selection = int(m.group(1)) - if 0 <= selection < len(recs): - return True, None, recs[selection] - else: - print(f"Invalid index {selection}") - return True, text, None - - elif m := match(r'^t (.*)', response): - print("searching for new matches") - text = m.group(1) - return True, text, None - - elif response.startswith("?"): - print(""" -Choices: -- Enter recommendation number to rename file, -- "t [text]" to search for new recommendations based on [text] -- "p" re-use the last selected cat-id -- "?" for this message -- "q" to quit -- or any other key to skip this file and continue to next file -""") - return True, text, None - elif response.startswith('q'): - return (False, None, None) - else: - return None - @click.group(epilog="For more information see " "") # @click.option('--verbose', flag_value='verbose', help='Verbose output') diff --git a/ucsinfer/recommend.py b/ucsinfer/recommend.py new file mode 100644 index 0000000..6c73c56 --- /dev/null +++ b/ucsinfer/recommend.py @@ -0,0 +1,62 @@ +# recommend.py + +from re import match + +from .inference import InferenceContext + +def print_recommendation(path: str | None, text: str, ctx: InferenceContext, + interactive_rename: bool): + """ + Print recommendations interactively. + + :returns: If interactive_rename is false or path is None, returns None. + If interactive, returns a tuple of bool, str | None, str | None where: + - if retval[0] is False, the user has requested processing quit. + - if retval[1] is a str, this is the text the user has entered to + perform a new inference instead of the file metadata. + `print_recommendation` should be called again with this argument. + - if retval[2] is a str, this is the catid the user has selected. + """ + recs = ctx.classify_text_ranked(text) + print("----------") + if path: + print(f"Path: {path}") + + print(f"Text: {text or ''}") + for i, r in enumerate(recs): + cat, subcat, _ = ctx.lookup_category(r) + print(f"- {i}: {r} ({cat}-{subcat})") + + if interactive_rename and path is not None: + response = input("#, t [text], ?, q > ") + + if m := match(r'^([0-9]+)', response): + selection = int(m.group(1)) + if 0 <= selection < len(recs): + return True, None, recs[selection] + else: + print(f"Invalid index {selection}") + return True, text, None + + elif m := match(r'^t (.*)', response): + print("searching for new matches") + text = m.group(1) + return True, text, None + + elif response.startswith("?"): + print(""" +Choices: +- Enter recommendation number to rename file, +- "t [text]" to search for new recommendations based on [text] +- "p" re-use the last selected cat-id +- "?" for this message +- "q" to quit +- or any other key to skip this file and continue to next file +""") + return True, text, None + elif response.startswith('q'): + return (False, None, None) + else: + return None + +