From c197e8458025180efcfbfda43b5d84decf372073 Mon Sep 17 00:00:00 2001
From: Iamhexi <igor@sosnowicz.eu>
Date: Fri, 4 Oct 2024 22:32:21 +0200
Subject: [PATCH] fix(main): fix choosing from material db in CLI

---
 knowledge_verificator/main.py       | 24 +++++++++++++++---------
 knowledge_verificator/utils/menu.py | 16 ++++++++++++----
 tests/test_qg.py                    |  2 ++
 3 files changed, 29 insertions(+), 13 deletions(-)

diff --git a/knowledge_verificator/main.py b/knowledge_verificator/main.py
index 79d9f46..0892f67 100755
--- a/knowledge_verificator/main.py
+++ b/knowledge_verificator/main.py
@@ -43,31 +43,37 @@
                 material_titles = [
                     material.title for material in material_db.materials
                 ]
-                material = choose_from_menu(material_titles, 'materials')
+                material = choose_from_menu(
+                    material_db.materials,
+                    plural_name='materials',
+                    attribute_to_show='title',
+                )
+
                 if material is None:
                     continue
 
-                paragraph = str(
+                PARAGRAPH = str(
                     choose_from_menu(material.paragraphs, 'paragraphs')
                 )
-                if paragraph is None:
+
+                if PARAGRAPH is None:
                     continue
 
                 console.print('Learn this paragraph: ')
-                console.print(paragraph)
+                console.print(PARAGRAPH)
                 console.print()
                 input('Press ENTER when ready.')
 
             case '2':
                 console.print('Enter a paragraph you would like to learn: ')
-                paragraph = input().strip()
+                PARAGRAPH = input().strip()
 
             case _:
                 console.print('Unrecognised option, try again!')
 
-        logger.debug('Loaded the following paragraph:\n %s', paragraph)
+        logger.debug('Loaded the following paragraph:\n %s', PARAGRAPH)
 
-        chosen_answer = chooser.choose_answer(paragraph=paragraph)
+        chosen_answer = chooser.choose_answer(paragraph=PARAGRAPH)
         if not chosen_answer:
             raise ValueError(
                 'The supplied paragaph is either too short or too general. '
@@ -83,7 +89,7 @@
         )
 
         question_with_context = qg_module.generate(
-            answer=chosen_answer, context=paragraph
+            answer=chosen_answer, context=PARAGRAPH
         )
         question = question_with_context['question']
         logger.debug(
@@ -97,7 +103,7 @@
 
         nli_module = NaturalLanguageInference()
         relation = nli_module.infer_relation(
-            premise=paragraph, hypothesis=user_answer
+            premise=PARAGRAPH, hypothesis=user_answer
         )
 
         match relation:
diff --git a/knowledge_verificator/utils/menu.py b/knowledge_verificator/utils/menu.py
index 73f30f1..05d6f6f 100644
--- a/knowledge_verificator/utils/menu.py
+++ b/knowledge_verificator/utils/menu.py
@@ -6,7 +6,10 @@
 
 
 def choose_from_menu(
-    menu_elements: list[str], plural_name: str, max_line_width: int = 40
+    menu_elements: list[Any],
+    plural_name: str,
+    attribute_to_show: str = '',
+    max_line_width: int = 40,
 ) -> Any | None:
     """
     Prompt a user to choose an element from a list via terminal.
@@ -15,15 +18,20 @@ def choose_from_menu(
         menu_elements (list[str]): List of elements to choose from.
             Elements should be convertible to `str` (implement `__str__` method).
         plural_name (str): Plural name of the elements. For example: options, paragraphs or names.
+        attribute_to_show (str): Attribute, which should be shown. If empty, print an entire object.
         max_line_width (int): Maximum line width in number of columns. By default: 40.
 
     Returns:
-        any | None: Element of a list or None if a user provided incorrect
-            value via a terminal.
+        any | None: Element of a list or None if a user provided incorrect value via a terminal.
     """
     console.print(f'Available {plural_name}:')
     for i, element in enumerate(menu_elements):
-        console.print(f'[{i+1}] {clip_text(element, max_line_width)}')
+        option_name = ''
+        if attribute_to_show:
+            option_name = getattr(element, attribute_to_show)
+        else:
+            option_name = element
+        console.print(f'[{i+1}] {clip_text(option_name, max_line_width)}')
     material_choice = input('Your choice: ')
     console.print()
 
diff --git a/tests/test_qg.py b/tests/test_qg.py
index 3fc6b18..ea9a45a 100644
--- a/tests/test_qg.py
+++ b/tests/test_qg.py
@@ -39,3 +39,5 @@ def test_basic_question_generation(
     }
 
     assert output == expected
+
+    # https://huggingface.co/sentence-transformers/all-distilroberta-v1