From 433078f7e3bec882fa5725ee66d2170dca96ac23 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Mon, 9 Mar 2026 21:10:06 +0000 Subject: [PATCH 1/5] Initial plan From bd79179f958e0a5a0c30c8def2f074ea0a323357 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Tue, 10 Mar 2026 05:10:28 +0000 Subject: [PATCH 2/5] fix: reduce Trie and newmm tokenizer memory usage MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - trie.py: lazy Node.children (None instead of {}); saves ~3.25 MiB (52 K leaf nodes × 64 B) for the default word-dict trie - trie.py: remove self.words set; add _word_count int; saves ~2 MiB of set-object overhead per Trie; rewrite __contains__, __iter__, __len__, add, remove to use trie traversal / DFS - trie.py: __iter__ uses recursive DFS with shared mutable prefix to avoid O(k^2) list copies - tcc_p.py: add tcc_pos_array() returning bytearray indexed by position; O(1) array lookup with len(text)+1 bytes vs. ~80 B/position in a set - newmm.py: use tcc_pos_array() instead of tcc_pos(); replace all four 'in valid_poss' set-membership tests with valid_poss[i] array indexing - tests: add _word_count accuracy tests (re-add, double-remove, iter); add tcc_pos_array edge-case tests (empty/None/int input, array length) Co-authored-by: bact <128572+bact@users.noreply.github.com> --- pythainlp/tokenize/newmm.py | 10 ++--- pythainlp/tokenize/tcc_p.py | 26 ++++++++++- pythainlp/util/trie.py | 87 ++++++++++++++++++++++++------------- tests/core/test_tokenize.py | 8 ++++ tests/core/test_util.py | 14 ++++++ 5 files changed, 110 insertions(+), 35 deletions(-) diff --git a/pythainlp/tokenize/newmm.py b/pythainlp/tokenize/newmm.py index b7db9acb0..d03e3c529 100644 --- a/pythainlp/tokenize/newmm.py +++ b/pythainlp/tokenize/newmm.py @@ -27,7 +27,7 @@ from pythainlp.util import Trie from pythainlp.tokenize import word_dict_trie -from pythainlp.tokenize.tcc_p import tcc_pos +from pythainlp.tokenize.tcc_p import tcc_pos_array # match non-Thai tokens # `|` is used as like "early return", @@ -86,7 +86,7 @@ def _onecut(text: str, custom_dict: Trie) -> Generator[str, None, None]: graph_size = 0 # keep track of graph size, if too big, force cutoff - valid_poss = tcc_pos(text) # breaking positions that are TCC-valid + valid_poss = tcc_pos_array(text) # bytearray of valid TCC break positions len_text = len(text) pos_list = [0] # priority queue of possible breaking positions @@ -95,7 +95,7 @@ def _onecut(text: str, custom_dict: Trie) -> Generator[str, None, None]: begin_pos = heappop(pos_list) for word in custom_dict.prefixes(text, begin_pos): end_pos_candidate = begin_pos + len(word) - if end_pos_candidate in valid_poss: + if valid_poss[end_pos_candidate]: graph[begin_pos].append(end_pos_candidate) graph_size = graph_size + 1 @@ -121,12 +121,12 @@ def _onecut(text: str, custom_dict: Trie) -> Generator[str, None, None]: end_pos = m.end() else: # Thai token, find minimum skip for pos in range(begin_pos + 1, len_text): - if pos in valid_poss: + if valid_poss[pos]: words = [ word for word in custom_dict.prefixes(text, pos) if ( - (pos + len(word) in valid_poss) + valid_poss[pos + len(word)] and not _PAT_THAI_TWOCHARS.match(word) ) ] diff --git a/pythainlp/tokenize/tcc_p.py b/pythainlp/tokenize/tcc_p.py index f631c36bc..d29f2f291 100644 --- a/pythainlp/tokenize/tcc_p.py +++ b/pythainlp/tokenize/tcc_p.py @@ -88,7 +88,7 @@ def tcc_pos(text: str) -> set[int]: """TCC positions :param str text: text to be tokenized into character clusters - :return: list of the ending position of subwords + :return: set of the ending positions of character clusters :rtype: set[int] """ if not text or not isinstance(text, str): @@ -103,6 +103,30 @@ def tcc_pos(text: str) -> set[int]: return p_set +def tcc_pos_array(text: str) -> bytearray: + """TCC positions as a bytearray. + + Returns a bytearray of length ``len(text) + 1`` where index ``i`` + is ``1`` if position ``i`` is a valid Thai Character Cluster boundary, + and ``0`` otherwise. Array-index lookup is faster and uses less + memory than set membership for large texts. + + :param str text: text to be tokenized into character clusters + :return: bytearray of valid TCC boundary flags, indexed by position + :rtype: bytearray + """ + if not text or not isinstance(text, str): + return bytearray(1) + + arr = bytearray(len(text) + 1) + p = 0 + for w in tcc(text): + p += len(w) + arr[p] = 1 + + return arr + + def segment(text: str) -> list[str]: """Subword segmentation diff --git a/pythainlp/util/trie.py b/pythainlp/util/trie.py index 8a4982708..836d12bf7 100644 --- a/pythainlp/util/trie.py +++ b/pythainlp/util/trie.py @@ -9,7 +9,7 @@ from __future__ import annotations from collections.abc import Iterable, Iterator -from typing import Union +from typing import Optional, Union class Trie(Iterable[str]): @@ -46,20 +46,20 @@ class Trie(Iterable[str]): # output: 5 """ - words: set[str] root: Node + _word_count: int class Node: __slots__: tuple[str, str] = ("end", "children") def __init__(self) -> None: self.end: bool = False - self.children: dict[str, Trie.Node] = {} + # Children dict is created on demand to reduce memory for leaf nodes. + self.children: Optional[dict[str, Trie.Node]] = None def __init__(self, words: Iterable[str]) -> None: - self.words: set[str] = set(words) + self._word_count: int = 0 self.root: Trie.Node = Trie.Node() - for word in words: self.add(word) @@ -67,43 +67,50 @@ def add(self, word: str) -> None: """Add a word to the trie. Spaces in front of and following the word will be removed. - :param str text: a word + :param str word: a word """ word = word.strip() - self.words.add(word) cur = self.root for ch in word: + if cur.children is None: + cur.children = {} child = cur.children.get(ch) - if not child: + if child is None: child = Trie.Node() cur.children[ch] = child cur = child - cur.end = True + if not cur.end: + cur.end = True + self._word_count += 1 def remove(self, word: str) -> None: """Remove a word from the trie. If the word is not found, do nothing. - :param str text: a word + :param str word: a word """ - # remove from set first - if word not in self.words: - return - self.words.remove(word) - # then remove from nodes - parent = self.root - data = [] # track path to leaf + # Navigate to the word's end node, recording the path. + node = self.root + path: list[tuple[Trie.Node, Trie.Node, str]] = [] for ch in word: - child = parent.children[ch] - data.append((parent, child, ch)) - parent = child - # remove the last one - child.end = False - # prune up the tree - for parent, child, ch in reversed(data): + if node.children is None: + return # word not in trie + child = node.children.get(ch) + if child is None: + return # word not in trie + path.append((node, child, ch)) + node = child + if not node.end: + return # path exists but not a complete word + node.end = False + self._word_count -= 1 + # Prune nodes that are now unused (not an end and no children). + for parent, child, ch in reversed(path): if child.end or child.children: break - del parent.children[ch] # remove from parent dict + del parent.children[ch] + if not parent.children: + parent.children = None # free empty dict def prefixes(self, text: str, start: int = 0) -> list[str]: """List all possible words from first sequence of characters in a word. @@ -118,8 +125,10 @@ def prefixes(self, text: str, start: int = 0) -> list[str]: i = start n = len(text) while i < n: + if cur.children is None: + break node = cur.children.get(text[i]) - if not node: + if node is None: break if node.end: res.append(text[start : i + 1]) @@ -128,13 +137,33 @@ def prefixes(self, text: str, start: int = 0) -> list[str]: return res def __contains__(self, key: str) -> bool: - return key in self.words + cur = self.root + for ch in key: + if cur.children is None: + return False + node = cur.children.get(ch) + if node is None: + return False + cur = node + return cur.end def __iter__(self) -> Iterator[str]: - yield from self.words + # DFS through the trie to yield all stored words. + # A shared mutable prefix list is appended/popped to avoid + # O(k²) list copies that a stack-based approach would incur. + def _dfs(node: Trie.Node, prefix: list[str]) -> Iterator[str]: + if node.end: + yield "".join(prefix) + if node.children: + for ch, child in node.children.items(): + prefix.append(ch) + yield from _dfs(child, prefix) + prefix.pop() + + yield from _dfs(self.root, []) def __len__(self) -> int: - return len(self.words) + return self._word_count def dict_trie(dict_source: Union[str, Iterable[str], Trie]) -> Trie: diff --git a/tests/core/test_tokenize.py b/tests/core/test_tokenize.py index a48617340..009c0bdf8 100644 --- a/tests/core/test_tokenize.py +++ b/tests/core/test_tokenize.py @@ -684,6 +684,14 @@ def test_tcc_p(self): # ) self.assertEqual(list(tcc_p.tcc("")), []) self.assertEqual(tcc_p.tcc_pos(""), set()) + # tcc_pos_array: edge cases + self.assertIsInstance(tcc_p.tcc_pos_array(""), bytearray) + self.assertIsInstance(tcc_p.tcc_pos_array(None), bytearray) + self.assertIsInstance(tcc_p.tcc_pos_array(42), bytearray) + # valid text: array length must equal len(text)+1 and mark boundaries + arr = tcc_p.tcc_pos_array("ประเทศ") + self.assertEqual(len(arr), len("ประเทศ") + 1) + self.assertEqual(arr[0], 0) # position 0 is never a boundary def test_display_cell_tokenize(self): self.assertEqual(display_cell_tokenize(""), []) diff --git a/tests/core/test_util.py b/tests/core/test_util.py index bd715bda9..3a02f49cf 100644 --- a/tests/core/test_util.py +++ b/tests/core/test_util.py @@ -516,6 +516,20 @@ def test_trie(self): trie.remove("ทด") self.assertEqual(len(trie), 2) + # _word_count must not double-count re-added words + trie2 = Trie(["ก", "ข", "ก"]) + self.assertEqual(len(trie2), 2) + trie2.add("ก") # already present – count must stay the same + self.assertEqual(len(trie2), 2) + trie2.add("ค") + self.assertEqual(len(trie2), 3) + trie2.remove("ข") + self.assertEqual(len(trie2), 2) + trie2.remove("ข") # removing non-existent word must not change count + self.assertEqual(len(trie2), 2) + # All remaining words must be reachable via __iter__ + self.assertEqual(sorted(trie2), ["ก", "ค"]) + trie = Trie([]) self.assertEqual(len(trie), 0) trie.remove("หมด") From f7add1e97ba1eb3512def60b789b47370e8d4722 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Tue, 10 Mar 2026 05:41:26 +0000 Subject: [PATCH 3/5] fix: resolve mypy union-attr error in Trie.remove(); add CHANGELOG entry Add explicit `if parent.children is not None:` guard before `del parent.children[ch]` in Trie.remove() so mypy can narrow the Optional[dict] type. The guard is always true at runtime but mypy cannot infer this from the path-traversal logic above. Add CHANGELOG.md entry for the Trie memory and tcc_pos_array changes. Co-authored-by: bact <128572+bact@users.noreply.github.com> --- CHANGELOG.md | 4 ++++ pythainlp/util/trie.py | 9 ++++++--- 2 files changed, 10 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 6739f39e9..7fb095875 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -94,6 +94,10 @@ The minimum requirement is now Python 3.9. downloaded (#1317) - `newmm` tokenization: Exponential-time explosion when text has many ambiguous breaking points (#1319) +- `Trie`: Reduce peak memory by using lazy `Node.children` allocation, + replacing `words: set[str]` with `_word_count: int`, and using + `tcc_pos_array()` bytearray in `newmm` for faster TCC boundary + lookups (#1323) ### Security diff --git a/pythainlp/util/trie.py b/pythainlp/util/trie.py index 836d12bf7..ee84289e7 100644 --- a/pythainlp/util/trie.py +++ b/pythainlp/util/trie.py @@ -105,12 +105,15 @@ def remove(self, word: str) -> None: node.end = False self._word_count -= 1 # Prune nodes that are now unused (not an end and no children). + # parent.children is always non-None here because the path was + # built by traversing through existing children dicts. for parent, child, ch in reversed(path): if child.end or child.children: break - del parent.children[ch] - if not parent.children: - parent.children = None # free empty dict + if parent.children is not None: # always true; narrows type + del parent.children[ch] + if not parent.children: + parent.children = None # free empty dict def prefixes(self, text: str, start: int = 0) -> list[str]: """List all possible words from first sequence of characters in a word. From 0af1b773967dc2ef977b66bbc347be46ecdf8a9b Mon Sep 17 00:00:00 2001 From: Arthit Suriyawongkul Date: Tue, 10 Mar 2026 13:04:35 +0700 Subject: [PATCH 4/5] Update CHANGELOG with Trie memory reduction summary Consolidate memory optimization details for Trie and TCC lookups. --- CHANGELOG.md | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 7fb095875..32c2fb024 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -94,10 +94,7 @@ The minimum requirement is now Python 3.9. downloaded (#1317) - `newmm` tokenization: Exponential-time explosion when text has many ambiguous breaking points (#1319) -- `Trie`: Reduce peak memory by using lazy `Node.children` allocation, - replacing `words: set[str]` with `_word_count: int`, and using - `tcc_pos_array()` bytearray in `newmm` for faster TCC boundary - lookups (#1323) +- `Trie`: Reduce memory usage and faster TCC lookups (#1323) ### Security From 7562b15f35ff5e48d31a0b6973b20e80656c885b Mon Sep 17 00:00:00 2001 From: Arthit Suriyawongkul Date: Tue, 10 Mar 2026 13:13:59 +0700 Subject: [PATCH 5/5] Update CHANGELOG for Trie memory usage improvement --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 32c2fb024..d9b591835 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -94,7 +94,7 @@ The minimum requirement is now Python 3.9. downloaded (#1317) - `newmm` tokenization: Exponential-time explosion when text has many ambiguous breaking points (#1319) -- `Trie`: Reduce memory usage and faster TCC lookups (#1323) +- `Trie`: Reduce memory usage and faster TCC boundary lookups (#1323) ### Security