diff --git a/pythainlp/tag/named_entity.py b/pythainlp/tag/named_entity.py index d3a0d303f..98b715207 100644 --- a/pythainlp/tag/named_entity.py +++ b/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 `_. + * *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." ) diff --git a/pythainlp/util/__init__.py b/pythainlp/util/__init__.py index eb65b23f4..068e51623 100644 --- a/pythainlp/util/__init__.py +++ b/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 ( diff --git a/tests/compact/testc_util.py b/tests/compact/testc_util.py index fa7b4cdf9..7ff9915e5 100644 --- a/tests/compact/testc_util.py +++ b/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): diff --git a/tests/core/test_tag.py b/tests/core/test_tag.py index 089fb2f6b..8bbfe309c 100644 --- a/tests/core/test_tag.py +++ b/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):