diff --git a/pythainlp/tag/named_entity.py b/pythainlp/tag/named_entity.py index 216ebe6a8..ab320a5cf 100644 --- a/pythainlp/tag/named_entity.py +++ b/pythainlp/tag/named_entity.py @@ -73,12 +73,22 @@ def _doc2features(doc, i) -> Dict: class ThaiNameTagger: - def __init__(self) -> None: + def __init__(self, version: str = "1.5") -> None: """ Thai named-entity recognizer. + + :param str version: Thai NER version. + It's support Thai NER 1.4 & 1.5. + The defualt value is `1.5` """ self.crf = CRFTagger() - self.crf.open(get_corpus_path(_CORPUS_NAME, version="1.5")) + + if version == "1.4": + self.crf.open(get_corpus_path("thainer-1.4", version="1.4")) + self.pos_tag_name = "orchid_ud" + else: + self.crf.open(get_corpus_path(_CORPUS_NAME, version="1.5")) + self.pos_tag_name = "lst20" def get_ner( self, text: str, pos: bool = True, tag: bool = False @@ -135,7 +145,11 @@ def get_ner( 'วันที่ 15 ก.ย. 61 ทดสอบระบบเวลา ' """ tokens = word_tokenize(text, engine=_TOKENIZER_ENGINE) - pos_tags = pos_tag(tokens, engine="perceptron", corpus="lst20") + pos_tags = pos_tag( + tokens, + engine="perceptron", + corpus=self.pos_tag_name + ) x_test = ThaiNameTagger.__extract_features(pos_tags) y = self.crf.tag(x_test) diff --git a/tests/test_tag.py b/tests/test_tag.py index f4ffef759..b70b1dc7b 100644 --- a/tests/test_tag.py +++ b/tests/test_tag.py @@ -206,6 +206,35 @@ def test_ner(self): "ทดสอบระบบเวลา ", ) + ner = ThaiNameTagger(version="1.4") + self.assertEqual(ner.get_ner(""), []) + self.assertIsNotNone(ner.get_ner("แมวทำอะไรตอนห้าโมงเช้า")) + self.assertIsNotNone(ner.get_ner("แมวทำอะไรตอนห้าโมงเช้า", pos=False)) + self.assertIsNotNone( + ner.get_ner( + """คณะวิทยาศาสตร์ประยุกต์และวิศวกรรมศาสตร์ มหาวิทยาลัยขอนแก่น + วิทยาเขตหนองคาย 112 หมู่ 7 บ้านหนองเดิ่น + ตำบลหนองกอมเกาะ อำเภอเมือง + จังหวัดหนองคาย 43000""" + ) + ) + self.assertIsNotNone( + ner.get_ner( + """คณะวิทยาศาสตร์ประยุกต์และวิศวกรรมศาสตร์ มหาวิทยาลัยขอนแก่น + วิทยาเขตหนองคาย 112 หมู่ 7 บ้านหนองเดิ่น + ตำบลหนองกอมเกาะ อำเภอเมือง + จังหวัดหนองคาย 43000""", + tag=True, + ) + ) + + # arguement `tag` is True + self.assertEqual( + ner.get_ner("วันที่ 15 ก.ย. 61 ทดสอบระบบเวลา 14:49 น.", tag=True), + "วันที่ 15 ก.ย. 61 " + "ทดสอบระบบเวลา ", + ) + self.assertEqual( ner.get_ner( "url = https://thainlp.org/pythainlp/docs/2.0/", tag=True @@ -260,7 +289,7 @@ def test_ner(self): # arguement `tag` is False and `pos` is True self.assertEqual( ner.get_ner("ไทย", pos=True, tag=False), - [("ไทย", "NN", "B-LOCATION")], + [('ไทย', 'PROPN', 'B-LOCATION')], ) # arguement `tag` is False and `pos` is False