diff --git a/ucsinfer/__main__.py b/ucsinfer/__main__.py index 1b5fab6..fcf5139 100644 --- a/ucsinfer/__main__.py +++ b/ucsinfer/__main__.py @@ -2,6 +2,8 @@ import os import sys import csv +from typing import Optional + from sentence_transformers import SentenceTransformer import tqdm import click @@ -11,19 +13,60 @@ from .inference import InferenceContext, load_ucs from .util import ffmpeg_description, parse_ucs -@click.group() +def recommend_text(text: str, ctx: InferenceContext): + return None + +@click.group(epilog="For more information see " + "") # @click.option('--verbose', flag_value='verbose', help='Verbose output') def ucsinfer(): + """ + Tools for applying UCS categories to sounds using large-language Models + + """ pass @ucsinfer.command('recommend') -def recommend(): +@click.option('--text', type=Optional[str], default=None, + help="Recommend a category for given text instead of reading " + "from a file") +@click.argument('files', nargs=-1) +@click.option('--model', type=str, metavar="", + default="paraphrase-multilingual-mpnet-base-v2", + show_default=True, + help="Select the sentence_transformer model to use") +def recommend(text, paths, model): """ Infer a UCS category for a text description """ - pass + m = SentenceTransformer(model) + ctx = InferenceContext(m, model) + + recommendations = [] + if text is not None: + recommendations.append({ + "text": text, + "recommendations": recommend_text(text, ctx) + }) + + for path in paths: + text = ffmpeg_description(path) + if text: + recommendations.append({ + "path":path, + "text":text, + "recommendations":recommend_text(text, ctx) + }) + else: + recommendations.append({ + "path":path, + "text":None, + "recommendations":None + }) + + @ucsinfer.command('gather')