diff --git a/fugashi/fugashi.pyx b/fugashi/fugashi.pyx index 597f538..ee5bba6 100644 --- a/fugashi/fugashi.pyx +++ b/fugashi/fugashi.pyx @@ -81,6 +81,10 @@ cdef class Node: def feature_raw(self): return self.c_node.feature.decode('utf-8') + @property + def id(self): + return self.c_node.id + @property def length(self): return self.c_node.length @@ -89,6 +93,14 @@ cdef class Node: def rlength(self): return self.c_node.rlength + @property + def rc_attr(self): + return self.c_node.rcAttr + + @property + def lc_attr(self): + return self.c_node.lcAttr + @property def posid(self): return self.c_node.posid @@ -101,6 +113,29 @@ cdef class Node: def stat(self): return self.c_node.stat + @property + def isbest(self): + # only relevant in nbest mode + return self.c_node.isbest + + @property + def alpha(self): + # only available if the tagger is started with the -m flag, else 0 + return self.c_node.alpha + + @property + def beta(self): + # only available if the tagger is started with the -m flag, else 0 + return self.c_node.beta + + @property + def wcost(self): + return self.c_node.wcost + + @property + def cost(self): + return self.c_node.cost + @property def is_unk(self): return self.stat == 1 @@ -250,7 +285,7 @@ cdef class GenericTagger: # This function just exists so subclasses can override the node type. return Node.wrap(node, self.wrapper) - def parseToNodeList(self, text): + def parseToNodeList(self, text, strip=True): # cstr = bytes(text, 'utf-8') bstr = bytes(text, 'utf-8') cdef const mecab_node_t* node = mecab_sparse_tonode(self.c_tagger, bstr) @@ -258,16 +293,13 @@ cdef class GenericTagger: # A nodelist always contains one each of BOS and EOS (beginning/end of # sentence) nodes. Since they have no information on them and MeCab # doesn't do any kind of sentence tokenization they're not useful in - # the output and will be removed here. + # the output and will be removed here by default (strip=True) # Node that on the command line this behavior is different, and each # line is treated as a sentence. out = [] - while node.next: - node = node.next - if node.stat == 3: # eos node - return out + while node: nn = self.wrap(node) # TODO maybe add an option to this function that doesn't cache the @@ -286,6 +318,17 @@ cdef class GenericTagger: nn.surface = self._cache[shash] out.append(nn) + node = node.next + + if strip: + # remove BOS and EOS nodes + out = out[1:-1] + else: + # set surface for BOS and EOS + out[0].surface = "BOS" + out[-1].surface = "EOS" + + return out def nbest(self, text, num=10): """Return the n-best possible tokenizations of the input, giving the @@ -296,7 +339,7 @@ cdef class GenericTagger: out = mecab_nbest_sparse_tostr(self.c_tagger, num, cstr).decode('utf-8') return out.rstrip() - def nbestToNodeList(self, text, num=10): + def nbestToNodeList(self, text, num=10, strip=True): """Return the n-best possible tokenizations of the input, giving each as a list of nodes. """ @@ -313,10 +356,7 @@ cdef class GenericTagger: # this happens if there aren't enough paths break out = [] - while node.next: - node = node.next - if node.stat == 3: - break + while node: nn = self.wrap(node) surf = node.surface[:node.length] shash = hash(surf) @@ -325,7 +365,16 @@ cdef class GenericTagger: self._cache[shash] = sys.intern(surf.decode("utf-8")) nn.surface = self._cache[shash] out.append(nn) + node = node.next + if strip: + # remove BOS and EOS nodes + out = out[1:-1] + else: + # set surface for BOS and EOS + out[0].surface = "BOS" + out[-1].surface = "EOS" + ret.append(out) return ret diff --git a/fugashi/mecab.pxd b/fugashi/mecab.pxd index 3ede409..65b2883 100644 --- a/fugashi/mecab.pxd +++ b/fugashi/mecab.pxd @@ -14,9 +14,16 @@ cdef extern from "mecab.h": unsigned int id unsigned short length unsigned short rlength + unsigned short rcAttr + unsigned short lcAttr unsigned short posid unsigned char char_type unsigned char stat + unsigned char isbest + float alpha + float beta + short wcost + long cost cdef struct mecab_model_t: pass diff --git a/fugashi/tests/test_basic.py b/fugashi/tests/test_basic.py index 9366bb4..6e2693c 100644 --- a/fugashi/tests/test_basic.py +++ b/fugashi/tests/test_basic.py @@ -86,3 +86,29 @@ def test_clobber(): nodes2 = tagger("x y z !") assert "a b c d".split() == [nn.surface for nn in nodes1] + +@pytest.mark.parametrize('text,wakati', WAKATI_TESTS) +def test_adding_bos_eos_nodes(text, wakati): + tagger = Tagger() + nodes1 = tagger.parseToNodeList(text, strip=True) + nodes2 = tagger.parseToNodeList(text, strip=False) + + assert len(nodes1) + 2 == len(nodes2) + assert nodes2[0].surface == "BOS" + assert nodes2[-1].surface == "EOS" + +def test_extended_attributes(): + tagger = Tagger("-m") + + nodes = tagger("ふがしは美味しい") + # we could test specific values of these, but for now this just tests that + # they are set and don't blow up + for node in nodes: + assert node.id is not None + assert node.rc_attr is not None + assert node.lc_attr is not None + assert node.wcost > 0 + assert node.cost > 0 + # these are not zero because we use the -m flag to turn on marginal probs + assert node.alpha != 0 + assert node.beta != 0