Merge pull request #1054 from PyThaiNLP/add-calculate_ngram_counts

Add pythainlp.lm.calculate_ngram_counts

Author: wannaphong

docs/api/lm.rst

@@ -6,4 +6,5 @@ pythainlp.lm
 Modules
 -------
 
+.. autofunction:: calculate_ngram_counts
 .. autofunction:: remove_repeated_ngrams

pythainlp/lm/__init__.py

@@ -3,6 +3,9 @@
 # SPDX-FileType: SOURCE
 # SPDX-License-Identifier: Apache-2.0
 
-__all__ = ["remove_repeated_ngrams"]
+__all__ = [
+    "calculate_ngram_counts",
+    "remove_repeated_ngrams"
+]
 
-from pythainlp.lm.text_util import remove_repeated_ngrams
+from pythainlp.lm.text_util import calculate_ngram_counts, remove_repeated_ngrams

pythainlp/lm/text_util.py

@@ -4,7 +4,32 @@
 # SPDX-License-Identifier: Apache-2.0
 # ruff: noqa: C901
 
-from typing import List
+from typing import List, Tuple, Dict
+
+
+def calculate_ngram_counts(
+        list_words: List[str],
+        n_min: int = 2,
+        n_max: int = 4) -> Dict[Tuple[str], int]:
+    """
+    Calculates the counts of n-grams in the list words for the specified range.
+
+    :param List[str] list_words: List of string
+    :param int n_min: The minimum n-gram size (default: 2).
+    :param int n_max: The maximum n-gram size (default: 4).
+
+    :return: A dictionary where keys are n-grams and values are their counts.
+    :rtype: Dict[Tuple[str], int]
+    """
+
+    ngram_counts = {}
+
+    for n in range(n_min, n_max + 1):
+        for i in range(len(list_words) - n + 1):
+            ngram = tuple(list_words[i:i + n])
+            ngram_counts[ngram] = ngram_counts.get(ngram, 0) + 1
+
+    return ngram_counts
 
 
 def remove_repeated_ngrams(string_list: List[str], n: int = 2) -> List[str]:

tests/core/test_lm.py

@@ -5,10 +5,23 @@
 
 import unittest
 
-from pythainlp.lm import remove_repeated_ngrams
+from pythainlp.lm import calculate_ngram_counts, remove_repeated_ngrams
 
 
 class LMTestCase(unittest.TestCase):
+    def test_calculate_ngram_counts(self):
+        self.assertEqual(
+            calculate_ngram_counts(['1', '2', '3', '4']),
+            {
+                ('1', '2'): 1,
+                ('2', '3'): 1,
+                ('3', '4'): 1,
+                ('1', '2', '3'): 1,
+                ('2', '3', '4'): 1,
+                ('1', '2', '3', '4'): 1
+            }
+        )
+
     def test_remove_repeated_ngrams(self):
         texts = ['เอา', 'เอา', 'แบบ', 'แบบ', 'แบบ', 'ไหน']
         self.assertEqual(