Add tag for dependency_parsing

Author: wannaphong

pythainlp/parse/core.py

@@ -1,16 +1,21 @@
 # -*- coding: utf-8 -*-
+from typing import List, Union
+
+
 _tagger = None
 _tagger_name = ""
 
-def dependency_parsing(text: str, model: str=None, engine: str="esupar")->str:
+def dependency_parsing(text: str, model: str=None, tag: str="str", engine: str="esupar")->Union[List[List[str]], str]:
     """
     Dependency Parsing
 
     :param str text: text to do dependency parsing
     :param str model: model for using with engine \
         (for esupar and transformers_ud)
+    :param str tag: output type (str or list)
     :param str engine: the name dependency parser
-    :return: str (conllu)
+    :return: str (conllu) or List
+    :rtype: Union[List[List[str]], str]
 
     **Options for engine**
         * *esupar* (default) - Tokenizer POS-tagger and Dependency-parser \
@@ -89,4 +94,4 @@ def dependency_parsing(text: str, model: str=None, engine: str="esupar")->str:
                 "The engine doesn't support."
             )
     _tagger_name = engine
-    return _tagger(text)
+    return _tagger(text, tag=tag)

pythainlp/parse/esupar_engine.py

@@ -4,6 +4,7 @@
 
 GitHub: https://github.com/KoichiYasuoka/esupar
 """
+from typing import List, Union
 import esupar
 
 
@@ -13,5 +14,12 @@ def __init__(self, model: str="th") -> None:
             model = "th"
         self.nlp=esupar.load(model)
 
-    def __call__(self, text):
-        return self.nlp(text)
+    def __call__(self, text: str, tag: str="str")->Union[List[List[str]], str]:
+        _data = str(self.nlp(text))
+        if tag =="list":
+            _temp = _data.splitlines()
+            _tag_data=[]
+            for i in _temp:
+                _tag_data.append(i.split())
+            return _tag_data
+        return _data

pythainlp/parse/spacy_thai_engine.py

@@ -4,16 +4,22 @@
 
 GitHub: https://github.com/KoichiYasuoka/spacy-thai
 """
+from typing import List, Union
 import spacy_thai
 
 
 class Parse:
     def __init__(self, model: str="th") -> None:
         self.nlp=spacy_thai.load()
 
-    def __call__(self, text:str)->str:
+    def __call__(self, text:str, tag: str="str")->Union[List[List[str]], str]:
         doc = self.nlp(text)
         _text = []
+        if tag == "list":
+            _tag_data=[]
+            for t in doc:
+                _tag_data.append([str(t.i+1),t.orth_,t.lemma_,t.pos_,t.tag_,"_",str(0 if t.head==t else t.head.i+1),t.dep_,"_","_" if t.whitespace_ else "SpaceAfter=No"])
+            return _tag_data
         for t in doc:
             _text.append("\t".join([str(t.i+1),t.orth_,t.lemma_,t.pos_,t.tag_,"_",str(0 if t.head==t else t.head.i+1),t.dep_,"_","_" if t.whitespace_ else "SpaceAfter=No"]))
         return '\n'.join(_text)

pythainlp/parse/transformers_ud.py

@@ -11,6 +11,7 @@
 GitHub: https://github.com/KoichiYasuoka
 """
 import os
+from typing import List, Union
 import numpy
 import torch
 import ufal.chu_liu_edmonds
@@ -48,7 +49,7 @@ def __init__(self, model: str="KoichiYasuoka/deberta-base-thai-ud-head") -> None
             tokenizer=self.tokenizer
         )
 
-    def __call__(self, text: str)->str:
+    def __call__(self, text: str, tag: str="str")->Union[List[List[str]], str]:
         w=[(t["start"],t["end"],t["entity_group"]) for t in self.deprel(text)]
         z,n={t["start"]:t["entity"].split("|") for t in self.tagger(text)},len(w)
         r,m=[text[s:e] for s,e,p in w],numpy.full((n+1,n+1),numpy.nan)
@@ -73,6 +74,12 @@ def __call__(self, text: str)->str:
             m[0:j,0]=m[j+1:,0]=numpy.nan
             h=ufal.chu_liu_edmonds.chu_liu_edmonds(m)[0]
         u=""
+        if tag == "list":
+            _tag_data=[]
+            for i,(s,e,p) in enumerate(w,1):
+                p="root" if h[i]==0 else "dep" if p=="root" else p
+                _tag_data.append([str(i),r[i-1],"_",z[s][0][2:],"_","|".join(z[s][1:]),str(h[i]),p,"_","_" if i<n and e<w[i][0] else "SpaceAfter=No"])
+            return _tag_data
         for i,(s,e,p) in enumerate(w,1):
             p="root" if h[i]==0 else "dep" if p=="root" else p
             u+="\t".join(

tests/test_parse.py

@@ -7,5 +7,8 @@
 class TestParsePackage(unittest.TestCase):
     def test_dependency_parsing(self):
         self.assertIsNotNone(dependency_parsing("ผมเป็นคนดี", engine="esupar"))
+        self.assertIsNotNone(dependency_parsing("ผมเป็นคนดี", engine="esupar", tag="list"))
         self.assertIsNotNone(dependency_parsing("ผมเป็นคนดี", engine="transformers_ud"))
+        self.assertIsNotNone(dependency_parsing("ผมเป็นคนดี", engine="transformers_ud", tag="list"))
         self.assertIsNotNone(dependency_parsing("ผมเป็นคนดี", engine="spacy_thai"))
+        self.assertIsNotNone(dependency_parsing("ผมเป็นคนดี", engine="spacy_thai", tag="list"))