perf(utils): EojeolCounter 멀티프로세싱 지원

soynlp/utils/utils.py

@@ -16,6 +16,21 @@
 installpath = os.path.sep.join(os.path.dirname(os.path.realpath(__file__)).split(os.path.sep)[:-1])
 
 
+def _count_eojeol_chunk(args: tuple) -> dict[str, int]:
+    """Worker function for parallel eojeol counting. Must be module-level for pickling."""
+    chunk, max_length, text_key, preprocess = args
+    counter: dict[str, int] = {}
+    for item in chunk:
+        sent = item[text_key] if isinstance(item, dict) else item
+        if preprocess is not None:
+            sent = preprocess(sent)
+        for eojeol in sent.split():
+            if (not eojeol) or (len(eojeol) > max_length):
+                continue
+            counter[eojeol] = counter.get(eojeol, 0) + 1
+    return counter
+
+
 def get_available_memory() -> float:
     """It returns remained memory as percentage"""
     mem = psutil.virtual_memory()
@@ -176,23 +191,34 @@ def __init__(
         verbose: bool = False,
         preprocess: Callable[[str], str] | None = None,
         text_key: str = "text",
+        n_workers: int = 1,
     ) -> None:
         self.min_count = min_count
         self.max_length = max_length
         self.filtering_checkpoint = filtering_checkpoint
         self.verbose = verbose
         self.text_key = text_key
 
+        self._has_custom_preprocess = preprocess is not None
         if preprocess is None:
 
             def base_preprocessing(x: str) -> str:
                 return x
 
             preprocess = base_preprocessing
+            self._parallel_preprocess: Callable[[str], str] | None = None
+        else:
+            import pickle
+
+            try:
+                pickle.dumps(preprocess)
+                self._parallel_preprocess = preprocess
+            except (pickle.PicklingError, AttributeError):
+                self._parallel_preprocess = None
         self.preprocess = preprocess
 
         if sents is not None:
-            self._counter = self._counting_from_sents(sents)
+            self._counter = self._counting_from_sents(sents, n_workers=n_workers)
         else:
             self._counter = {}
 
@@ -209,13 +235,17 @@ def __getitem__(self, eojeol: str) -> int:
     def __len__(self) -> int:
         return len(self._counter)
 
-    def _counting_from_sents(self, sents: Any) -> dict[str, int]:
+    def _counting_from_sents(self, sents: Any, n_workers: int = 1) -> dict[str, int]:
         check_corpus(sents)
+        if n_workers != 1 and (not self._has_custom_preprocess or self._parallel_preprocess is not None):
+            return self._counting_from_sents_parallel(sents, n_workers, self._parallel_preprocess)
+        if n_workers != 1:
+            logger.info("EojeolCounter: custom preprocess가 pickle 불가 — 단일 프로세스로 집계")
         if self.verbose:
             sent_iterator = tqdm(sents, desc="[EojeolCounter] counting eojeols ", total=len(sents))
         else:
             sent_iterator = sents
-        counter = {}
+        counter: dict[str, int] = {}
         for i_sent, item in enumerate(sent_iterator):
             if isinstance(item, dict):
                 sent = item[self.text_key]
@@ -231,6 +261,26 @@ def _counting_from_sents(self, sents: Any) -> dict[str, int]:
         counter = {eojeol: count for eojeol, count in counter.items() if count >= self.min_count}
         return counter
 
+    def _counting_from_sents_parallel(
+        self, sents: Any, n_workers: int, preprocess: Callable[[str], str] | None = None
+    ) -> dict[str, int]:
+        from multiprocessing import Pool, cpu_count
+
+        texts = list(sents)
+        n = cpu_count() if n_workers == -1 else n_workers
+        chunk_size = max(1, len(texts) // n)
+        chunks = [texts[i : i + chunk_size] for i in range(0, len(texts), chunk_size)]
+        worker_args = [(chunk, self.max_length, self.text_key, preprocess) for chunk in chunks]
+
+        with Pool(processes=n) as pool:
+            partial_counters = pool.map(_count_eojeol_chunk, worker_args)
+
+        merged: dict[str, int] = {}
+        for partial in partial_counters:
+            for eojeol, count in partial.items():
+                merged[eojeol] = merged.get(eojeol, 0) + count
+        return {eojeol: count for eojeol, count in merged.items() if count >= self.min_count}
+
     def remove_eojeols(self, eojeols: set[str] | str) -> "EojeolCounter":
         """Remove eojeols
 

tests/unit/test_utils.py

@@ -235,3 +235,35 @@ def test_corpus_loader_with_eojeol_counter(jsonl_file):
     counter = EojeolCounter(sents=loader)
     assert counter["이것은"] == 1
     assert counter["예문입니다"] == 1
+
+
+# --- EojeolCounter multiprocessing tests ---
+
+_MULTI_SENTS = [
+    "아이오아이가 평가단에게 높은 점수를 받았습니다",
+    "아이오아이는 아이돌 그룹입니다",
+    "자연어처리는 어렵습니다",
+    "자연어처리를 공부합니다",
+    "명사추출은 중요한 작업입니다",
+] * 100
+
+
+def test_eojeol_counter_multi_equals_single():
+    """n_workers=4 결과가 단일 프로세스 결과와 동일하다."""
+    single = EojeolCounter(sents=_MULTI_SENTS, n_workers=1)
+    multi = EojeolCounter(sents=_MULTI_SENTS, n_workers=4)
+    assert dict(single.items()) == dict(multi.items())
+
+
+def test_eojeol_counter_auto_workers():
+    """n_workers=-1이면 CPU 코어 수를 자동 사용한다."""
+    single = EojeolCounter(sents=_MULTI_SENTS, n_workers=1)
+    auto = EojeolCounter(sents=_MULTI_SENTS, n_workers=-1)
+    assert dict(single.items()) == dict(auto.items())
+
+
+def test_eojeol_counter_multi_with_min_count():
+    """n_workers > 1일 때 min_count 필터가 올바르게 적용된다."""
+    single = EojeolCounter(sents=_MULTI_SENTS, min_count=5, n_workers=1)
+    multi = EojeolCounter(sents=_MULTI_SENTS, min_count=5, n_workers=4)
+    assert dict(single.items()) == dict(multi.items())