Merge pull request #1255 from bact/add-more-tests

Add more NER

Author: bact

pythainlp/tag/named_entity.py

@@ -16,12 +16,15 @@ class NER:
     :param str corpus: corpus
 
     **Options for engine**
-        * *thainer-v2* - Thai NER engine v2.0 for Thai NER 2.0 (default)
+        * *phayathaibert* - PhayaThaiBERT-based Thai NER engine
         * *thainer* - Thai NER engine
+        * *thainer-v2* - Thai NER engine v2.0 for Thai NER 2.0 (default)
         * *tltk* - wrapper for `TLTK <https://pypi.org/project/tltk/>`_.
+        * *wangchanberta* - WangchanBERTa-based Thai NER engine
 
     **Options for corpus**
         * *thainer* - Thai NER corpus (default)
+        * *thainer-v2* - Thai NER v2 corpus
 
     **Note**: The tltk engine supports NER models from tltk only.
     """
@@ -34,29 +37,33 @@ def __init__(
     def load_engine(self, engine: str, corpus: str) -> None:
         self.name_engine = engine
         self.engine: Any = None
-        if engine == "thainer" and corpus == "thainer":
-            from pythainlp.tag.thainer import ThaiNameTagger
+        if corpus == "thainer":
+            if engine == "thainer":
+                from pythainlp.tag.thainer import ThaiNameTagger
 
-            self.engine = ThaiNameTagger()
-        elif engine == "thainer-v2" and corpus == "thainer":
-            from pythainlp.wangchanberta import NamedEntityRecognition
+                self.engine = ThaiNameTagger()
+            elif engine == "thainer-v2":
+                from pythainlp.wangchanberta import NamedEntityRecognition
 
-            self.engine = NamedEntityRecognition(
-                model="pythainlp/thainer-corpus-v2-base-model"
-            )
-        elif engine == "tltk":
-            from pythainlp.tag import tltk
+                self.engine = NamedEntityRecognition(
+                    model="pythainlp/thainer-corpus-v2-base-model"
+                )
+            elif engine == "wangchanberta":
+                from pythainlp.wangchanberta import ThaiNameTagger as WangchanbertaThaiNameTagger  # type: ignore[assignment]  # noqa: I001,E501
+
+                self.engine = WangchanbertaThaiNameTagger(dataset_name=corpus)  # type: ignore[call-arg]
+        elif corpus == "thainer-v2":
+            if engine == "phayathaibert":
+                from pythainlp.phayathaibert.core import NamedEntityTagger
 
-            self.engine = tltk
-        elif engine == "wangchanberta" and corpus == "thainer":
-            from pythainlp.wangchanberta import ThaiNameTagger  # type: ignore[assignment]  # noqa: I001
+                self.engine = NamedEntityTagger()
+        else:  # No corpus matched
+            if engine == "tltk":
+                from pythainlp.tag import tltk
 
-            self.engine = ThaiNameTagger(dataset_name=corpus)  # type: ignore[call-arg]
-        elif engine == "phayathaibert" and corpus == "thainer-v2":
-            from pythainlp.phayathaibert.core import NamedEntityTagger
+                self.engine = tltk
 
-            self.engine = NamedEntityTagger()
-        else:
+        if self.engine is None:
             raise ValueError(
                 f"NER class not support {engine} engine or {corpus} corpus."
             )

pythainlp/util/__init__.py

@@ -45,6 +45,8 @@
     "reorder_vowels",
     "rhyme",
     "sound_syllable",
+    "spell_syllable",
+    "spell_word",
     "spelling",
     "spell_words",
     "syllable_length",
@@ -121,6 +123,7 @@
 from pythainlp.util.remove_trailing_repeat_consonants import (
     remove_trailing_repeat_consonants,
 )
+from pythainlp.util.spell_words import spell_syllable, spell_word
 from pythainlp.util.strftime import thai_strftime
 
 from pythainlp.util.thai import (

tests/compact/testc_util.py

@@ -6,8 +6,7 @@
 
 import unittest
 
-from pythainlp.util import rhyme, thai_word_tone_detector
-from pythainlp.util.spell_words import spell_word
+from pythainlp.util import rhyme, spell_word, thai_word_tone_detector
 
 
 class SpellWordTestCaseC(unittest.TestCase):

tests/core/test_tag.py

@@ -90,9 +90,12 @@ def test_pos_tag(self):
         )
 
     def test_NER_error_handling(self):
-        # Test error handling for invalid engine/corpus combination
         with self.assertRaises(ValueError):
-            NER(engine="thainer", corpus="cat")
+            NER(engine="xx_non_existing", corpus="thainer")
+        with self.assertRaises(ValueError):
+            NER(engine="xx_non_existing", corpus="thainer-v2")
+        with self.assertRaises(ValueError):
+            NER(engine="xx_non_existing", corpus="xx_non_existing")
 
 
 class PerceptronTaggerTestCase(unittest.TestCase):