Add NERCut

Author: wannaphong

pythainlp/tokenize/core.py

@@ -162,6 +162,10 @@ def word_tokenize(
     elif engine == "icu":
         from .pyicu import segment
 
+        segments = segment(text)
+    elif engine == "nercut":
+        from .nercut import segment
+
         segments = segment(text)
     else:
         raise ValueError(

pythainlp/tokenize/nercut.py

@@ -0,0 +1,66 @@
+# -*- coding: utf-8 -*-
+"""
+nercut 0.1
+
+Code by Wannaphong Phatthiyaphaibun
+"""
+from typing import List
+from pythainlp.tag.named_entity import ThaiNameTagger
+
+_thainer = ThaiNameTagger()
+
+def segment(
+    text: str,
+    tag:List[str] = [
+        "ORGANIZATION",
+        "PERSON",
+        "PHONE",
+        "EMAIL",
+        "DATE",
+        "TIME"
+    ]
+) -> List[str]:
+    """
+    nercut 0.1
+
+    Code by Wannaphong Phatthiyaphaibun
+
+    neww+thainer word segmentation.
+
+    :param str text: text to be tokenized to words
+    :parm list tag: ThaiNER tag
+    :return: list of words, tokenized from the text
+    """
+    global _thainer
+    if not text or not isinstance(text, str):
+        return []
+
+    _ws = _thainer.get_ner(text, pos = False)
+    _list_w = []
+    _bi = ""
+    _tag = ""
+    for i,t in _ws:
+        if t != "O":
+            _tag_temp = t.split('-')[1]
+        else:
+            _tag_temp = "O"
+        if t.startswith('B-') and _tag_temp in tag:
+            if _bi!="" and _tag in tag:
+                _list_w.append(_bi)
+            _bi=""
+            _bi += i
+            _tag = t.replace('B-','')
+        elif t.startswith('I-') and t.replace('I-','') == _tag and _tag_temp in tag:
+            _bi += i
+        elif t == "O" and _tag != "" and _tag in tag:
+            _list_w.append(_bi)
+            _bi=""
+            _tag = ""
+            _list_w.append(i)
+        else:
+            _bi=""
+            _tag = ""
+            _list_w.append(i)
+    if _bi!="":
+        _list_w.append(_bi)
+    return _list_w

tests/test_tokenize.py

@@ -20,6 +20,7 @@
 )
 from pythainlp.tokenize.ssg import segment as ssg_segment
 from pythainlp.util import dict_trie
+from pythainlp.tokenize import nercut as tokenize_nercut
 
 
 class TestTokenizePackage(unittest.TestCase):
@@ -230,6 +231,7 @@ def test_word_tokenize(self):
         self.assertIsNotNone(word_tokenize(self.text_1, engine="icu"))
         self.assertIsNotNone(word_tokenize(self.text_1, engine="deepcut"))
         self.assertIsNotNone(word_tokenize(self.text_1, engine="attacut"))
+        self.assertIsNotNone(word_tokenize(self.text_1, engine="nercut"))
         with self.assertRaises(ValueError):
             word_tokenize("หมอนทอง", engine="XX")  # engine does not exist
 
@@ -364,6 +366,15 @@ def test_word_tokenize_attacut(self):
             ["ฉัน", "รัก", "ภาษา", "ไทย", "เพราะ", "ฉัน", "เป็น", "คน", "ไทย"],
         )
 
+    def test_word_tokenize_nercut(self):
+        self.assertEqual(tokenize_nercut.segment(None), [])
+        self.assertEqual(tokenize_nercut.segment(""), [])
+        self.assertIsNotNone(
+            tokenize_nercut.segment("ทดสอบ")
+        )
+        self.assertIsNotNone(tokenize_nercut.segment("ทดสอบ"))
+        self.assertIsNotNone(word_tokenize("ทดสอบ", engine="nercut"))
+
     def test_sent_tokenize(self):
         self.assertEqual(sent_tokenize(None), [])
         self.assertEqual(sent_tokenize(""), [])