diff --git a/ucsinfer/__main__.py b/ucsinfer/__main__.py index 6fcd340..5316b99 100644 --- a/ucsinfer/__main__.py +++ b/ucsinfer/__main__.py @@ -97,8 +97,8 @@ def recommend_category(path, embeddings, model) -> Tuple[str, list]: return desc, classify_text_ranked(desc, embeddings, model) -def lookup_cat(catid: str, ucs: list) -> tuple[str,str]: - return next( ((x['Category'], x['SubCategory']) \ +def lookup_cat(catid: str, ucs: list) -> tuple[str,str, str]: + return next( ((x['Category'], x['SubCategory'], x['Explanations']) \ for x in ucs if x['CatID'] == catid)) @@ -110,18 +110,18 @@ class Commands(cmd.Cmd): self.file_list = [] self.model: Optional[SentenceTransformer] = None self.embeddings: Optional[list] = None - self.catlist: Optional[list] = None - self.model_rec_list = [] + self.catlist: list = [] + self.rec_list = [] self.history = [] def default(self, line): try: rec = int(line) - if rec < len(self.model_rec_list): + if rec < len(self.rec_list): print(f"Accept option {rec}") - ind = rec - 1 - self.history = [self.model_rec_list[ind]] + self.history[0:4] + ind = rec + self.history = [self.rec_list[ind]] + self.history[0:4] self.onecmd("next") else: pass @@ -129,42 +129,63 @@ class Commands(cmd.Cmd): except ValueError: super().default(line) - def preloop(self) -> None: self.file_cursor = 0 self.update_prompt() + self.setup_for_file() return super().preloop() + def precmd(self, line: str): + + return super().precmd(line) + def postcmd(self, stop: bool, line: str) -> bool: - self.update_prompt() + if not stop: + self.update_prompt() + self.setup_for_file() return super().postcmd(stop, line) def update_prompt(self): self.prompt = f"(ucsinfer:{self.file_cursor}/{len(self.file_list)}) " - + + def setup_for_file(self): + file = self.file_list[self.file_cursor] + desc, recs = recommend_category(file, self.embeddings, self.model) + self.onecmd('file') + print(f" >> {desc}") + self.print_recommendations(recs) + + def print_recommendations(self, top_recs): + + def print_one_rec(index, rec): + cat, subcat, exp = lookup_cat(rec, self.catlist) + print(f" [{index:2}] : {rec} - {cat} / {subcat} - {exp[0:30]}") + + self.rec_list = [] + print("Suggested from description:") + for rec in top_recs: + print_one_rec(len(self.rec_list), rec) + self.rec_list.append(rec) + + if len(self.history) > 0: + print("History:") + for rec in self.history: + print_one_rec(len(self.rec_list), rec) + self.rec_list.append(rec) + def do_file(self, _): 'Print info about the current file' + print("---") if self.file_cursor < len(self.file_list): - self.update_prompt() path = self.file_list[self.file_cursor] f = os.path.basename(path) - print(f" > {f}") + print(f" > {f}") else: - print( " > No file") - - def do_rec(self, _): - if self.file_cursor < len(self.file_list): - self.update_prompt() - path = self.file_list[self.file_cursor] - desc, recs = recommend_category(path, self.embeddings, self.model) - print(f" >> {desc}") - self.model_rec_list = recs - else: - self.model_rec_list = [] + print( " > No file") def do_lookup(self, args): 'print a list of UCS categories similar to the argument' - self.model_rec_list = classify_text_ranked(args, self.embeddings, + self.rec_list = classify_text_ranked(args, self.embeddings, self.model) def do_ls(self, _): @@ -172,17 +193,19 @@ class Commands(cmd.Cmd): for file in self.file_list[self.file_cursor:] + \ self.file_list[0:self.file_cursor]: f = os.path.basename(file) - print(f" > {f}") + print(f" > {f}") def do_next(self, _): 'go to next file' self.file_cursor += 1 self.file_cursor = self.file_cursor % len(self.file_list) + self.setup_for_file() def do_prev(self, _): 'go to previous file' self.file_cursor -= 1 self.file_cursor = self.file_cursor % len(self.file_list) + self.setup_for_file() def do_bye(self, _): 'exit the program'