feat: add backend support for multi-word stemming with test cases

my_ghost_writer/constants.py

@@ -17,5 +17,6 @@ LOG_JSON_FORMAT = bool(os.getenv("LOG_JSON_FORMAT"))
 IS_TESTING = bool(os.getenv('IS_TESTING', ""))
 LOG_LEVEL = os.getenv("LOG_LEVEL", "INFO")
 API_MODE = bool(os.getenv("API_MODE", ""))
+N_WORDS_GRAM = int(os.getenv("N_WORDS_GRAM", 2))
 session_logger.setup_logging(json_logs=LOG_JSON_FORMAT, log_level=LOG_LEVEL)
 app_logger = structlog.stdlib.get_logger(__name__)

my_ghost_writer/text_parsers.py

@@ -2,14 +2,14 @@ from typing import Iterator
 
 from nltk import PorterStemmer
 
-from my_ghost_writer.constants import app_logger
+from my_ghost_writer.constants import app_logger, N_WORDS_GRAM
 from my_ghost_writer.type_hints import RequestTextRowsParentList, ResponseTextRowsDict
 
 
 ps = PorterStemmer()
 
 
-def text_stemming(text: str | RequestTextRowsParentList) -> ResponseTextRowsDict:
+def text_stemming(text: str | RequestTextRowsParentList, n = 3) -> ResponseTextRowsDict:
     """
     Applies Porter Stemmer algorithm to reduce words in a given text to their base form;
     then it uses WordPunctTokenizer() to produce a dict of words frequency with, for
@@ -17,6 +17,7 @@ def text_stemming(text: str | RequestTextRowsParentList) -> ResponseTextRowsDict
 
     Args:
         text (str): Input string containing the text to be stemmed.
+        n (int): The maximum number of words to consider for n-grams (default is 3).
 
     Returns:
         tuple[int, dict]: a tuple with the number of processed total rows within the initial text and the word frequency dict
@@ -43,9 +44,12 @@ def text_stemming(text: str | RequestTextRowsParentList) -> ResponseTextRowsDict
     idx_rows = []
     idx_rows_child = []
     idx_rows_parent = []
+    rows_dict = {}
     for textrow in valid_textrows_with_num:
         row = textrow["text"]
-        idx_rows.append(textrow["idxRow"])
+        idx_row = textrow["idxRow"]
+        rows_dict[idx_row] = row
+        idx_rows.append(idx_row)
         try:
             idx_rows_child.append(textrow["idxRowChild"])
             idx_rows_parent.append(textrow["idxRowParent"])
@@ -54,39 +58,11 @@ def text_stemming(text: str | RequestTextRowsParentList) -> ResponseTextRowsDict
             idx_rows_parent.append(None)
         row_words_tokens.append(wordpunct_tokenize(row))
         row_offsets_tokens.append(WordPunctTokenizer().span_tokenize(row))
-    words_stems_dict = get_words_tokens_and_indexes(row_words_tokens, row_offsets_tokens, idx_rows, idx_rows_child, idx_rows_parent)
+    words_stems_dict = get_words_tokens_and_indexes_ngrams(row_words_tokens, row_offsets_tokens, idx_rows, idx_rows_child, idx_rows_parent, rows_dict=rows_dict, n=n)
     n_total_rows = len(valid_textrows_with_num)
     return n_total_rows, words_stems_dict
 
 
-def get_words_tokens_and_indexes(
-        words_tokens_list: list[str], offsets_tokens_list: list | Iterator, idx_rows_list: list[int], idx_rows_child: list[int], idx_rows_parent: list[int]
-    ) -> dict:
-    """
-    Get the word tokens and their indexes in the text.
-
-    Args:
-        words_tokens_list (list): List of words tokens.
-        offsets_tokens_list (list): List of offsets for each token.
-        idx_rows_list (list[int]): List of row indices corresponding to the tokens.
-        idx_rows_child (list[int]): List of child row indices corresponding to the tokens.
-        idx_rows_parent (list[int]): List of parent row indices corresponding to the tokens.
-
-    Returns:
-        dict: Dictionary with stemmed words as keys and a list of dictionaries
-              containing the original word and its offsets as values.
-    """
-    words_stems_dict = {}
-    for (n_row, n_row_child, n_row_parent, words_tokens, offsets_tokens) in zip(idx_rows_list, idx_rows_child, idx_rows_parent, words_tokens_list, offsets_tokens_list):
-        for word, offsets in zip(words_tokens, offsets_tokens):
-            stem = ps.stem(word)
-            if stem not in words_stems_dict:
-                words_stems_dict[stem] = {"count": 0, "word_prefix": stem, "offsets_array": []}
-            count, word_offsets = update_stems_list(words_stems_dict[stem], word, offsets, n_row=n_row, n_row_child=n_row_child, n_row_parent=n_row_parent)
-            words_stems_dict[stem] = {"count": count, "word_prefix": stem, "offsets_array": word_offsets}
-    return words_stems_dict
-
-
 def update_stems_list(current_stem_tuple: dict, word: str, offsets: list, n_row: int, n_row_child: int, n_row_parent: int) -> tuple:
     """
     Update the stem list with the new stem and its count.
@@ -106,3 +82,73 @@ def update_stems_list(current_stem_tuple: dict, word: str, offsets: list, n_row:
     n += 1
     word_offsets.append({"word": word, "offsets": list(offsets), "n_row": n_row, "n_row_child": n_row_child, "n_row_parent": n_row_parent})
     return n, word_offsets
+
+
+def get_words_tokens_and_indexes_ngrams(
+        words_tokens_list: list[list[str]] | Iterator,
+        offsets_tokens_list: list[list[tuple[int, int]]] | Iterator,
+        idx_rows_list: list[int],
+        idx_rows_child: list[int],
+        idx_rows_parent: list[int],
+        rows_dict: dict[int, str],
+        n: int = N_WORDS_GRAM
+) -> dict:
+    f"""
+    Like get_words_tokens_and_indexes, but supports joined n-grams (from 1 up to n words).
+    Returns a dict with n-gram stem as key and offsets/count as in example_result.
+    The 'word_prefix' is set to the most common 'word' in offsets_array.
+
+    Args:
+        words_tokens_list (list): List of lists of words tokens.
+        offsets_tokens_list (list): List of lists of offsets for each token.
+        idx_rows_list (list[int]): List of row indices corresponding to the tokens.
+        idx_rows_child (list[int]): List of child row indices corresponding to the tokens.
+        idx_rows_parent (list[int]): List of parent row indices corresponding to the tokens.
+        rows_dict (dict[int, str]): Dictionary mapping row indices to their text.
+        n (int): The maximum number of words to consider for n-grams (default is from the N_WORDS_GRAM constant,
+                 right now it has value of ${N_WORDS_GRAM}).
+
+    Returns:
+        dict: Dictionary with n-gram stems as keys and a dictionary of their counts, word prefixes, and offsets as values.
+    """
+    from collections import Counter
+
+    ngram_dict = {}
+    for (n_row, n_row_child, n_row_parent, words_tokens, offsets_tokens) in zip(
+            idx_rows_list, idx_rows_child, idx_rows_parent, words_tokens_list, offsets_tokens_list
+    ):
+        words_tokens = list(words_tokens)
+        offsets_tokens = list(offsets_tokens)
+        length = len(words_tokens)
+        for n_words_ngram in range(1, n + 1):
+            for i in range(length - n_words_ngram + 1):
+                row = rows_dict[n_row]
+                ngram_words = words_tokens[i:i + n_words_ngram]
+                stem_list = [ps.stem(word=word) for word in ngram_words]
+                ngram_offsets = offsets_tokens[i:i + n_words_ngram]
+                start = ngram_offsets[0][0]
+                end = ngram_offsets[-1][1]
+                ngram_stem = " ".join(stem_list)
+                ngram = row[start:end]
+                if ngram_stem not in ngram_dict:
+                    ngram_dict[ngram_stem] = {"count": 0, "word_prefix": ngram, "offsets_array": [], "n_words_ngram": n_words_ngram}
+                # Use update_stems_list to update count and offsets_array
+                count, offsets_array = update_stems_list(
+                    ngram_dict[ngram_stem],
+                    ngram,
+                    [start, end],
+                    n_row=n_row,
+                    n_row_child=n_row_child,
+                    n_row_parent=n_row_parent
+                )
+                ngram_dict[ngram_stem]["count"] = count
+                ngram_dict[ngram_stem]["offsets_array"] = offsets_array
+
+    # Update word_prefix to the most common 'word' in offsets_array
+    for entry in ngram_dict.values():
+        words = [item["word"] for item in entry["offsets_array"] if "word" in item]
+        if words:
+            most_common_word, _ = Counter(words).most_common(1)[0]
+            entry["word_prefix"] = most_common_word
+
+    return ngram_dict

tests/events/llm_generated_story_4.txt

@@ -0,0 +1,3 @@
+Once upon a time, in the kingdom of Elwoodia, there lived a young and wise young Princess named Isabella. She saw a tiny creature like one of the fairies her mother told her.
+One of the fairies flew to her and spoke in a soft voice. The fairy was dressed with colourful flowers and green leaves, with a colourful flower with many petals.
+The fairy grabbed some rocks. One was red, crowned with colourful flowers and some green leaves. Many of the fairies were small but quick.

tests/test_text_parsers.py

@@ -7,6 +7,29 @@ from nltk.tokenize import wordpunct_tokenize, WordPunctTokenizer
 from tests import EVENTS_FOLDER
 
 
+def get_inputs_for(valid_textrows_with_num):
+    row_words_tokens = []
+    row_offsets_tokens = []
+    idx_rows = []
+    idx_rows_child = []
+    idx_rows_parent = []
+    rows_dict = {}
+    for textrow in valid_textrows_with_num:
+        row = textrow["text"]
+        idx_row = textrow["idxRow"]
+        rows_dict[idx_row] = row
+        idx_rows.append(idx_row)
+        try:
+            idx_rows_child.append(textrow["idxRowChild"])
+            idx_rows_parent.append(textrow["idxRowParent"])
+        except KeyError:
+            idx_rows_child.append(None)
+            idx_rows_parent.append(None)
+        row_words_tokens.append(wordpunct_tokenize(row))
+        row_offsets_tokens.append(WordPunctTokenizer().span_tokenize(row))
+    return row_words_tokens, row_offsets_tokens, idx_rows, idx_rows_child, idx_rows_parent, rows_dict
+
+
 class TestTextParsers(unittest.TestCase):
     def setUp(self):
         with open(EVENTS_FOLDER / "get_words_tokens_and_indexes_inputs.json", "r") as src:
@@ -34,35 +57,17 @@ class TestTextParsers(unittest.TestCase):
             self.ps = nltk.PorterStemmer()
             self.wnl = nltk.WordNetLemmatizer()
 
-    def test_get_words_tokens_and_indexes(self):
-        from my_ghost_writer.text_parsers import get_words_tokens_and_indexes
-        row_words_tokens = self.get_words_tokens_and_indexes_inputs["row_words_tokens"]
-        row_offsets_tokens = self.get_words_tokens_and_indexes_inputs["row_offsets_tokens"]
-        idx_rows = self.get_words_tokens_and_indexes_inputs["idx_rows"]
-        idx_rows_child = self.get_words_tokens_and_indexes_inputs["idx_rows_child"]
-        idx_rows_parent = self.get_words_tokens_and_indexes_inputs["idx_rows_parent"]
-        words_stems_dict = get_words_tokens_and_indexes(
-            row_words_tokens,
-            row_offsets_tokens,
-            idx_rows,
-            idx_rows_child,
-            idx_rows_parent
-        )
-        with open(EVENTS_FOLDER / "response_text_stemming_with_parents.json", "r") as dst_json:
-            response_text_stemming_no_parents = json.load(dst_json)
-            expected_words_stems_dict = response_text_stemming_no_parents["words_stems_dict"]
-        self.assertDictEqual(words_stems_dict, expected_words_stems_dict)
-
     def test_text_stemming_text(self):
         from my_ghost_writer.text_parsers import text_stemming
         self.maxDiff = None
         n_total_rows, words_stems_dict = text_stemming(self.original_text)
         self.assertEqual(n_total_rows, len(self.text_split_newline))
         # with open(EVENTS_FOLDER / "response_text_stemming_from_llm_generated_story_3.json", "w") as dst_json:
-        #     json.dump(words_stems_dict, dst_json, indent=2)
+        #     json.dump({"n_total_rows": n_total_rows, "words_stems_dict": words_stems_dict}, dst_json, indent=2)
+        #     pass
         with open(EVENTS_FOLDER / "response_text_stemming_from_llm_generated_story_3.json", "r") as dst_json:
-            response_text_stemming_no_parents = json.load(dst_json)
-            expected_words_stems_dict = response_text_stemming_no_parents["words_stems_dict"]
+            response_text_stemming_from_llm_generated_story_3 = json.load(dst_json)
+            expected_words_stems_dict = response_text_stemming_from_llm_generated_story_3["words_stems_dict"]
         self.assertDictEqual(words_stems_dict, expected_words_stems_dict)
 
     def test_text_stemming_input_str_json(self):
@@ -72,11 +77,11 @@ class TestTextParsers(unittest.TestCase):
         n_total_rows, words_stems_dict = text_stemming(json_str)
         self.assertEqual(n_total_rows, len(self.text_json_list_no_parents))
         # with open(EVENTS_FOLDER / "response_text_stemming_empty_rows.json", "w") as dst_json:
-        #     json.dump(words_stems_dict, dst_json, indent=2)
+        #     json.dump({"n_total_rows": n_total_rows, "words_stems_dict": words_stems_dict}, dst_json, indent=2)
         # pass
-        with open(EVENTS_FOLDER / "response_text_stemming_no_parents.json", "r") as dst_json:
-            response_text_stemming_no_parents = json.load(dst_json)
-            expected_words_stems_dict = response_text_stemming_no_parents["words_stems_dict"]
+        with open(EVENTS_FOLDER / "response_text_stemming_empty_rows.json", "r") as dst_json:
+            response_text_stemming_empty_rows = json.load(dst_json)
+            expected_words_stems_dict = response_text_stemming_empty_rows["words_stems_dict"]
         self.assertDictEqual(words_stems_dict, expected_words_stems_dict)
 
     def test_text_stemming_list_no_parents(self):
@@ -84,7 +89,7 @@ class TestTextParsers(unittest.TestCase):
         self.maxDiff = None
         n_total_rows, words_stems_dict = text_stemming(self.text_json_list_no_parents)
         self.assertEqual(n_total_rows, len(self.text_json_list_no_parents))
-        # with open(EVENTS_FOLDER / "response_text_stemming_with_parents2.json", "w") as dst_json:
+        # with open(EVENTS_FOLDER / "response_text_stemming_no_parents.json", "w") as dst_json:
         #     json.dump({"n_total_rows": n_total_rows, "words_stems_dict": words_stems_dict}, dst_json, indent=2)
         # pass
         with open(EVENTS_FOLDER / "response_text_stemming_no_parents.json", "r") as dst_json:
@@ -95,9 +100,9 @@ class TestTextParsers(unittest.TestCase):
     def test_text_stemming_list_with_parents(self):
         from my_ghost_writer.text_parsers import text_stemming
         self.maxDiff = None
-        n_total_rows, words_stems_dict = text_stemming(self.text_json_list_with_parents)
+        n_total_rows, words_stems_dict = text_stemming(self.text_json_list_with_parents, n=3)
         self.assertEqual(n_total_rows, len(self.text_json_list_with_parents))
-        # with open(EVENTS_FOLDER / "response_text_stemming_with_parents2.json", "w") as dst_json:
+        # with open(EVENTS_FOLDER / "response_text_stemming_with_parents.json", "w") as dst_json:
         #     json.dump({"n_total_rows": n_total_rows, "words_stems_dict": words_stems_dict}, dst_json, indent=2)
         # pass
         with open(EVENTS_FOLDER / "response_text_stemming_with_parents.json", "r") as dst_json:
@@ -134,5 +139,52 @@ class TestTextParsers(unittest.TestCase):
             self.assertEqual(word_offsets, expected_offsets_array)
 
 
+    def test_get_words_tokens_and_indexes_ngrams_no_parents(self):
+        from my_ghost_writer.text_parsers import get_words_tokens_and_indexes_ngrams
+
+        with open(EVENTS_FOLDER / "llm_generated_story_4.txt", "r") as src:
+            text = src.read()
+            valid_textrows_with_num = [{"idxRow": i, "text": row} for i, row in enumerate(text.split("\n"))]
+
+        row_words_tokens, row_offsets_tokens, idx_rows, idx_rows_child, idx_rows_parent, rows_dict = get_inputs_for(
+            valid_textrows_with_num
+        )
+        words_stems_dict = get_words_tokens_and_indexes_ngrams(
+            row_words_tokens,
+            row_offsets_tokens,
+            idx_rows,
+            idx_rows_child,
+            idx_rows_parent,
+            rows_dict=rows_dict,
+            n=5
+        )
+        # with open(EVENTS_FOLDER / "response_get_words_tokens_and_indexes_ngrams_text4_n5.json", "w") as dst_json:
+        #     json.dump({"words_stems_dict": words_stems_dict}, dst_json, indent=2)
+        with open(EVENTS_FOLDER / "response_get_words_tokens_and_indexes_ngrams_text4_n5.json", "r") as dst_json:
+            response_get_words_tokens_and_indexes_ngrams_text4_n5 = json.load(dst_json)
+            expected_words_stems_dict = response_get_words_tokens_and_indexes_ngrams_text4_n5["words_stems_dict"]
+        self.assertDictEqual(words_stems_dict, expected_words_stems_dict)
+
+    def test_get_words_tokens_and_indexes_ngrams_with_parents(self):
+        from my_ghost_writer.text_parsers import get_words_tokens_and_indexes_ngrams
+        self.maxDiff = None
+        row_words_tokens, row_offsets_tokens, idx_rows, idx_rows_child, idx_rows_parent, rows_dict = get_inputs_for(
+            self.text_json_list_with_parents
+        )
+        words_stems_dict = get_words_tokens_and_indexes_ngrams(
+            row_words_tokens,
+            row_offsets_tokens,
+            idx_rows,
+            idx_rows_child,
+            idx_rows_parent,
+            rows_dict=rows_dict,
+            n=3
+        )
+        with open(EVENTS_FOLDER / "response_text_stemming_with_parents.json", "r") as dst_json:
+            response_text_stemming_with_parents = json.load(dst_json)
+            expected_words_stems_dict = response_text_stemming_with_parents["words_stems_dict"]
+        self.assertDictEqual(words_stems_dict, expected_words_stems_dict)
+
+
 if __name__ == "__main__":
     unittest.main()