refactor: refactor NLTKHelper

Author: ChenZiHong-Gavin

graphgen/utils/help_nltk.py

@@ -1,39 +1,54 @@
+from functools import lru_cache
 import os
-from typing import Dict, List, Optional
+from typing import Dict, List, Final, Optional
 import nltk
 import jieba
 
-resource_path = os.path.join(os.path.dirname(os.path.dirname(__file__)), "resources")
-
-
 class NLTKHelper:
-    _stopwords: Dict[str, Optional[List[str]]] = {
-        "english": None,
-        "chinese": None,
+    """
+    NLTK helper class
+    """
+
+    SUPPORTED_LANGUAGES: Final[Dict[str, str]] = {
+        "en": "english",
+        "zh": "chinese"
+    }
+    _NLTK_PACKAGES: Final[Dict[str, str]] = {
+        "stopwords": "corpora",
+        "punkt_tab": "tokenizers"
     }
 
-    def __init__(self):
+    def __init__(self, nltk_data_path: Optional[str] = None):
+        self._nltk_path = nltk_data_path or os.path.join(
+            os.path.dirname(os.path.dirname(__file__)),
+            "resources", 
+            "nltk_data"
+        )
+        nltk.data.path.append(self._nltk_path)
         jieba.initialize()
 
+        self._ensure_nltk_data("stopwords")
+        self._ensure_nltk_data("punkt_tab")
+
+    def _ensure_nltk_data(self, package_name: str) -> None:
+        """
+        ensure nltk data is downloaded
+        """
+        try:
+            nltk.data.find(f"{self._NLTK_PACKAGES[package_name]}/{package_name}")
+        except LookupError:
+            nltk.download(package_name, download_dir=self._nltk_path, quiet=True)
+
+    @lru_cache(maxsize=2)
     def get_stopwords(self, lang: str) -> List[str]:
-        nltk.data.path.append(os.path.join(resource_path, "nltk_data"))
-        if self._stopwords[lang] is None:
-            try:
-                nltk.data.find("corpora/stopwords")
-            except LookupError:
-                nltk.download("stopwords", download_dir=os.path.join(resource_path, "nltk_data"))
-
-            self._stopwords[lang] = nltk.corpus.stopwords.words(lang)
-        return self._stopwords[lang]
-
-    @staticmethod
-    def word_tokenize(text: str, lang: str) -> List[str]:
+        if lang not in self.SUPPORTED_LANGUAGES:
+            raise ValueError(f"Language {lang} is not supported.")
+        return nltk.corpus.stopwords.words(self.SUPPORTED_LANGUAGES[lang])
+
+    def word_tokenize(self, text: str, lang: str) -> List[str]:
+        if lang not in self.SUPPORTED_LANGUAGES:
+            raise ValueError(f"Language {lang} is not supported.")
         if lang == "zh":
             return jieba.lcut(text)
-        nltk.data.path.append(os.path.join(resource_path, "nltk_data"))
-        try:
-            nltk.data.find("tokenizers/punkt_tab")
-        except LookupError:
-            nltk.download("punkt_tab", download_dir=os.path.join(resource_path, "nltk_data"))
 
         return nltk.word_tokenize(text)